Forward batch_size and seq_length to the validate_model_outputs in order to test with potential different input shapes.

mfuntowicz · mfuntowicz · commit db3b72dc215e · 2021-08-31T11:25:30.000+02:00
Signed-off-by: Morgan Funtowicz &lt;funtowiczmo@gmail.com&gt;
diff --git a/src/transformers/onnx/convert.py b/src/transformers/onnx/convert.py
@@ -138,19 +138,27 @@ def validate_model_outputs(
     reference_model: Union[PreTrainedModel, TFPreTrainedModel],
     onnx_model_path_or_bytes: Union[PathLike, bytes],
     onnx_named_outputs: List[str],
-    atol: float,
+    batch_size: int = -1,
+    seq_length: int = -1,
+    atol: float = 1e-5,
 ):
     from onnxruntime import InferenceSession, SessionOptions
 
     logger.info("Validating ONNX model...")
 
-    # TODO: generate inputs with a different batch_size and seq_len that was used for conversion to properly test
-    # dynamic input shapes.
-    reference_model_inputs = config.generate_dummy_inputs(tokenizer, framework=TensorType.PYTORCH)
-
     # Create ONNX Runtime session
     options = SessionOptions()
-    session = InferenceSession(onnx_model.as_posix(), options)
+    options.add_session_config_entry('session.load_model_format', 'ONNX')
+    session = InferenceSession(onnx_model_path_or_bytes, options)
+
+    # TODO: generate inputs with a different batch_size and seq_len that was used for conversion to properly test
+    # dynamic input shapes.
+    reference_model_inputs = config.generate_dummy_inputs(
+        tokenizer=tokenizer,
+        batch_size=batch_size if batch_size > 0 else 3,
+        seq_length=seq_length if seq_length > 0 else 31,
+        framework=TensorType.PYTORCH
+    )
 
     # Compute outputs from the reference model
     ref_outputs = reference_model(**reference_model_inputs)
diff --git a/tests/test_onnx_v2.py b/tests/test_onnx_v2.py
@@ -237,13 +237,25 @@ def test_pytorch_export_default(self):
                 model = model_class(config_class.from_pretrained(model))
                 onnx_config = onnx_config_class.from_model_config(model.config)
 
-                with NamedTemporaryFile("w") as output:
+                with NamedTemporaryFile("wb+") as output:
                     onnx_inputs, onnx_outputs = export(
-                        tokenizer, model, onnx_config, DEFAULT_ONNX_OPSET, Path(output.name)
+                        tokenizer, model, onnx_config, DEFAULT_ONNX_OPSET, output
                     )
 
                     try:
-                        validate_model_outputs(onnx_config, tokenizer, model, Path(output.name), onnx_outputs, 1e-5)
+                        # Reset to the head of the file and read everything
+                        output.seek(0)
+                        model_bytes = output.read()
+                        validate_model_outputs(
+                            onnx_config,
+                            tokenizer,
+                            model,
+                            model_bytes,
+                            onnx_outputs,
+                            batch_size=-1,
+                            seq_length=-1,
+                            atol=1e-5
+                        )
                     except ValueError as ve:
                         self.fail(f"{name} -> {ve}")
 
@@ -265,11 +277,22 @@ def test_pytorch_export_with_past(self):
                     onnx_config.use_past, "OnnxConfigWithPast.use_past should be if called with with_past()"
                 )
 
-                with NamedTemporaryFile("w") as output:
-                    output = Path(output.name)
+                with NamedTemporaryFile("wb+") as output:
                     onnx_inputs, onnx_outputs = export(tokenizer, model, onnx_config, DEFAULT_ONNX_OPSET, output)
 
                     try:
-                        validate_model_outputs(onnx_config, tokenizer, model, output, onnx_outputs, 1e-5)
+                        # Reset to the head of the file and read everything
+                        output.seek(0)
+                        model_bytes = output.read()
+                        validate_model_outputs(
+                            onnx_config,
+                            tokenizer,
+                            model,
+                            model_bytes,
+                            onnx_outputs,
+                            batch_size=-1,
+                            seq_length=-1,
+                            atol=1e-5
+                        )
                     except ValueError as ve:
                         self.fail(f"{name} -> {ve}")
