Merge pull request #687 from take-cheeze/pfto_test_names

[onnx] Output pfto tests to out/ dir like export_testcase

Author: emcastillo

pytorch_pfn_extras/onnx/export_testcase.py

@@ -328,6 +328,7 @@ def export_testcase(
     if isinstance(outs, torch.Tensor):
         outs = outs,
     assert outs is not None
+    outs = torch._C._jit_flatten(outs)[0]
     # Remove unused inputs
     # - When keep_initializers_as_inputs=True, inputs contains initializers.
     #   So we have to filt initializers.

tests/pytorch_pfn_extras_tests/onnx_tests/test_export_testcase.py

@@ -22,6 +22,7 @@
 
 
 output_dir = 'out'
+output_counter = {}
 
 
 class Net(nn.Module):
@@ -47,14 +48,24 @@ def _get_output_dir(d, **kwargs):
     output_dir_base = 'out'
     opset_ver = kwargs.get('opset_version', pytorch_pfn_extras.onnx._constants.onnx_default_opset)
 
+    test_name = os.getenv("PYTEST_CURRENT_TEST").split(':')[-1].split(' ')[0]
+    if d:
+        test_name = f"{test_name}_{d}"
+    if "model_overwrite" not in kwargs:
+        if test_name in output_counter:
+            output_counter[test_name] += 1
+            test_name = f"{test_name}_{output_counter[test_name]}"
+        else:
+            output_counter[test_name] = 0
+
     output_dir = os.path.join(
-        output_dir_base, 'opset{}'.format(opset_ver), d)
+        output_dir_base, 'opset{}'.format(opset_ver), test_name)
     os.makedirs(output_dir, exist_ok=True)
     return output_dir
 
 
 def _helper(model, args, d, use_pfto=True, **kwargs):
-    output_dir = _get_output_dir(d)
+    output_dir = _get_output_dir(d, **kwargs)
     if 'training' not in kwargs:
         kwargs['training'] = model.training
     if 'do_constant_folding' not in kwargs:

tests/pytorch_pfn_extras_tests/onnx_tests/utils.py

@@ -1,10 +1,12 @@
-import tempfile
-from typing import Callable, List, Optional
+import os
+from typing import Callable
 
 import onnx
 import onnxruntime as ort
 import torch
 import pytorch_pfn_extras.onnx.pfto_exporter.export as pfto
+from pytorch_pfn_extras.onnx import export_testcase
+from pytorch_pfn_extras_tests.onnx_tests.test_export_testcase import _get_output_dir
 
 
 def run_model_test(
@@ -13,8 +15,6 @@ def run_model_test(
     check_torch_export=True,
     rtol=1e-05,
     atol=1e-08,
-    input_names: Optional[List[str]] = None,
-    output_names: Optional[List[str]] = None,
     skip_oxrt=False,
     operator_export_type=torch.onnx.OperatorExportTypes.ONNX,
     strict_trace=True,
@@ -35,63 +35,60 @@ def run_model_test(
 
     if operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN:
         skip_oxrt = True
-    with tempfile.NamedTemporaryFile() as f:
-        f.close()
-        rng_state = torch.get_rng_state()
-        with pfto._force_tracing():
-            expected = model(*args)
-        if not isinstance(expected, tuple):
-            expected = (expected,)
 
-        te_model = None
-        if check_torch_export:
-            torch.set_rng_state(rng_state)
-            with tempfile.NamedTemporaryFile() as torch_f:
-                torch_f.close()
-                torch.onnx.export(
-                    model,
-                    args,
-                    torch_f.name,
-                    input_names=input_names,
-                    output_names=output_names,
-                    **kwargs,
-                )
-                te_model = onnx.load(torch_f.name)
+    rng_state = torch.get_rng_state()
+    with pfto._force_tracing():
+        expected = model(*args)
+    if isinstance(expected, torch.Tensor):
+        expected = (expected,)
+    expected = torch._C._jit_flatten(expected)[0]
 
-        if input_names is None:
-            input_names = [f"input_{idx}" for idx, _ in enumerate(args)]
-        if output_names is None:
-            output_names = [f"output_{idx}" for idx, _ in enumerate(expected)]
+    te_model = None
+    if check_torch_export:
         torch.set_rng_state(rng_state)
-        actual = pfto.export(
+        pt_dir = _get_output_dir("pt", **kwargs)
+        export_testcase(
             model,
             args,
-            f.name,
-            input_names=input_names,
-            output_names=output_names,
-            strict_trace=strict_trace,
+            pt_dir,
+            use_pfto=False,
             **kwargs,
         )
-        if not isinstance(actual, tuple):
-            actual = (actual,)
-        assert len(actual) == len(expected)
+        te_model = onnx.load(os.path.join(pt_dir, "model.onnx"))
 
-        for a, e in zip(actual, expected):
-            if isinstance(a, torch.Tensor) and isinstance(e, torch.Tensor):
-                assert torch.isclose(a, e, rtol=rtol, atol=atol).all()
+    torch.set_rng_state(rng_state)
+    pf_dir = _get_output_dir("pf", **kwargs)
+    actual = export_testcase(
+        model,
+        args,
+        pf_dir,
+        strict_trace=strict_trace,
+        return_output=True,
+        use_pfto=True,
+        **kwargs,
+    )
+    if isinstance(actual, torch.Tensor):
+        actual = (actual,)
+    expected = torch._C._jit_flatten(expected)[0]
+    assert len(actual) == len(expected)
 
-        pfto_model = onnx.load(f.name)
-        if te_model is not None:
-            assert len(te_model.graph.output) == len(pfto_model.graph.output)
-            assert len(te_model.graph.input) == len(pfto_model.graph.input)
+    for a, e in zip(actual, expected):
+        if isinstance(a, torch.Tensor) and isinstance(e, torch.Tensor):
+            assert torch.isclose(a, e, rtol=rtol, atol=atol).all()
 
-        if skip_oxrt:
-            return pfto_model
-
-        ort_session = ort.InferenceSession(f.name)
-        actual = ort_session.run(None, {k: v.cpu().numpy() for k, v in zip(input_names, args)})
-        for a, e in zip(actual, expected):
-            cmp = torch.isclose(torch.tensor(a), e.cpu(), rtol=rtol, atol=atol)
-            assert cmp.all(), f"{cmp.logical_not().count_nonzero()} / {cmp.numel()} values failed"
+    pfto_model = onnx.load(os.path.join(pf_dir, "model.onnx"))
+    if te_model is not None:
+        assert len(te_model.graph.output) == len(pfto_model.graph.output)
+        assert len(te_model.graph.input) == len(pfto_model.graph.input)
 
+    if skip_oxrt:
         return pfto_model
+
+    ort_session = ort.InferenceSession(os.path.join(pf_dir, "model.onnx"))
+    input_names = [i.name for i in pfto_model.graph.input]
+    actual = ort_session.run(None, {k: v.cpu().numpy() for k, v in zip(input_names, args)})
+    for a, e in zip(actual, expected):
+        cmp = torch.isclose(torch.tensor(a), e.cpu(), rtol=rtol, atol=atol)
+        assert cmp.all(), f"{cmp.logical_not().count_nonzero()} / {cmp.numel()} values failed"
+
+    return pfto_model