[ONNX] Create common ExportTestCase base class (pytorch#88145)

BowenBao · pytorchmergebot · commit 500fd65531e7 · 2022-11-10T21:51:59.000Z
Refactor out a common base class `ExportTestCase`, for common things in `setUp`. Pull Request resolved: pytorch#88145 Approved by: https://github.com/justinchuby, https://github.com/abock, https://github.com/AllenTiTaiWang
diff --git a/test/onnx/onnx_test_common.py b/test/onnx/onnx_test_common.py
@@ -3,15 +3,13 @@
 from __future__ import annotations
 
 import os
-import random
 from typing import Any, Mapping, Type
 
-import numpy as np
 import onnxruntime
+import pytorch_test_common
 
 import torch
 from torch.onnx import _constants, verification
-from torch.testing._internal import common_utils
 
 onnx_model_dir = os.path.join(
     os.path.dirname(os.path.realpath(__file__)),
@@ -54,21 +52,15 @@ def parameterize_class_name(cls: Type, idx: int, input_dicts: Mapping[Any, Any])
     return f"{cls.__name__}_{suffix}"
 
 
-def set_rng_seed(seed):
-    torch.manual_seed(seed)
-    random.seed(seed)
-    np.random.seed(seed)
-
-
-class _TestONNXRuntime(common_utils.TestCase):
+class _TestONNXRuntime(pytorch_test_common.ExportTestCase):
     opset_version = _constants.ONNX_DEFAULT_OPSET
     keep_initializers_as_inputs = True  # For IR version 3 type export.
     is_script = False
     check_shape = True
     check_dtype = True
 
     def setUp(self):
-        set_rng_seed(0)
+        super().setUp()
         onnxruntime.set_seed(0)
         if torch.cuda.is_available():
             torch.cuda.manual_seed_all(0)
diff --git a/test/onnx/pytorch_test_common.py b/test/onnx/pytorch_test_common.py
@@ -2,12 +2,17 @@
 
 import functools
 import os
+import random
 import sys
 import unittest
 from typing import Optional
 
+import numpy as np
+
 import torch
 from torch.autograd import function
+from torch.onnx._internal import diagnostics
+from torch.testing._internal import common_utils
 
 pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
 sys.path.insert(-1, pytorch_test_dir)
@@ -188,3 +193,24 @@ def wrapper(self, *args, **kwargs):
 
 def flatten(x):
     return tuple(function._iter_filter(lambda o: isinstance(o, torch.Tensor))(x))
+
+
+def set_rng_seed(seed):
+    torch.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+
+class ExportTestCase(common_utils.TestCase):
+    """Test case for ONNX export.
+
+    Any test case that tests functionalities under torch.onnx should inherit from this class.
+    """
+
+    def setUp(self):
+        super().setUp()
+        # TODO(#88264): Flaky test failures after changing seed.
+        set_rng_seed(0)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(0)
+        diagnostics.engine.clear()
diff --git a/test/onnx/test_autograd_funs.py b/test/onnx/test_autograd_funs.py
@@ -1,16 +1,16 @@
 # Owner(s): ["module: onnx"]
 
-import unittest
+import pytorch_test_common
 
 import torch
-
 from onnx_test_common import run_model_test
 from torch.onnx import OperatorExportTypes
 from torch.onnx._globals import GLOBALS
 from torch.onnx.utils import _model_to_graph
+from torch.testing._internal import common_utils
 
 
-class TestAutogradFuns(unittest.TestCase):
+class TestAutogradFuns(pytorch_test_common.ExportTestCase):
     opset_version = GLOBALS.export_onnx_opset_version
     keep_initializers_as_inputs = False
     onnx_shape_inference = True
@@ -209,4 +209,4 @@ def forward(self, input):
 
 
 if __name__ == "__main__":
-    unittest.main()
+    common_utils.run_tests()
diff --git a/test/onnx/test_custom_ops.py b/test/onnx/test_custom_ops.py
@@ -4,14 +4,15 @@
 import numpy as np
 import onnx
 import onnx_test_common
+import pytorch_test_common
 import torch
 import torch.utils.cpp_extension
 from test_pytorch_onnx_caffe2 import do_export
 from torch.onnx import symbolic_helper
 from torch.testing._internal import common_utils
 
 
-class TestCustomOps(common_utils.TestCase):
+class TestCustomOps(pytorch_test_common.ExportTestCase):
     def test_custom_add(self):
         op_source = """
         #include <torch/script.h>
@@ -56,7 +57,7 @@ def symbolic_custom_add(g, self, other):
         np.testing.assert_array_equal(caffe2_out[0], model(x, y).cpu().numpy())
 
 
-class TestCustomAutogradFunction(common_utils.TestCase):
+class TestCustomAutogradFunction(pytorch_test_common.ExportTestCase):
     opset_version = 9
     keep_initializers_as_inputs = False
     onnx_shape_inference = True
@@ -130,7 +131,7 @@ def symbolic_pythonop(ctx: torch.onnx.SymbolicContext, g, *args, **kwargs):
         onnx_test_common.run_model_test(self, model, input_args=(x,))
 
 
-class TestExportAsContribOps(common_utils.TestCase):
+class TestExportAsContribOps(pytorch_test_common.ExportTestCase):
     opset_version = 14
     keep_initializers_as_inputs = False
     onnx_shape_inference = True
diff --git a/test/onnx/test_export_modes.py b/test/onnx/test_export_modes.py
@@ -15,11 +15,13 @@
 # Make the helper files in test/ importable
 pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
 sys.path.append(pytorch_test_dir)
+import pytorch_test_common
+
 from torch.testing._internal import common_utils
 
 
 # Smoke tests for export methods
-class TestExportModes(common_utils.TestCase):
+class TestExportModes(pytorch_test_common.ExportTestCase):
     class MyModel(nn.Module):
         def __init__(self):
             super(TestExportModes.MyModel, self).__init__()
diff --git a/test/onnx/test_models.py b/test/onnx/test_models.py
@@ -2,8 +2,9 @@
 
 import unittest
 
-import torch
+import pytorch_test_common
 
+import torch
 from model_defs.dcgan import _netD, _netG, bsz, imgsz, nz, weights_init
 from model_defs.emb_seq import EmbeddingNetwork1, EmbeddingNetwork2
 from model_defs.mnist import MNIST
@@ -44,7 +45,7 @@ def toC(x):
 BATCH_SIZE = 2
 
 
-class TestModels(common_utils.TestCase):
+class TestModels(pytorch_test_common.ExportTestCase):
     opset_version = 9  # Caffe2 doesn't support the default.
     keep_initializers_as_inputs = False
 
diff --git a/test/onnx/test_models_onnxruntime.py b/test/onnx/test_models_onnxruntime.py
@@ -8,6 +8,7 @@
 import onnx_test_common
 import parameterized
 import PIL
+import pytorch_test_common
 import test_models
 
 import torch
@@ -64,7 +65,7 @@ def exportTest(
 
 TestModels = type(
     "TestModels",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(
         test_models.TestModels.__dict__,
         is_script_test_enabled=False,
@@ -77,7 +78,7 @@ def exportTest(
 # model tests for scripting with new JIT APIs and shape inference
 TestModels_new_jit_API = type(
     "TestModels_new_jit_API",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(
         TestModels.__dict__,
         exportTest=exportTest,
diff --git a/test/onnx/test_onnx_opset.py b/test/onnx/test_onnx_opset.py
@@ -4,6 +4,7 @@
 import itertools
 
 import onnx
+import pytorch_test_common
 
 import torch
 import torch.onnx
@@ -70,7 +71,7 @@ def check_onnx_opsets_operator(
         check_onnx_opset_operator(model, ops[opset_version], opset_version)
 
 
-class TestONNXOpset(common_utils.TestCase):
+class TestONNXOpset(pytorch_test_common.ExportTestCase):
     def test_opset_fallback(self):
         class MyModule(Module):
             def forward(self, x):
diff --git a/test/onnx/test_operators.py b/test/onnx/test_operators.py
@@ -31,6 +31,7 @@
 )
 from torch.autograd import Function, Variable
 from torch.nn import functional, Module
+from torch.onnx._internal import diagnostics
 from torch.onnx.symbolic_helper import (
     _get_tensor_dim_size,
     _get_tensor_sizes,
@@ -71,6 +72,10 @@ def forward(self, *args):
 
 
 class TestOperators(common_utils.TestCase):
+    def setUp(self):
+        super().setUp()
+        diagnostics.engine.clear()
+
     def assertONNX(self, f, args, params=None, **kwargs):
         if params is None:
             params = ()
diff --git a/test/onnx/test_pytorch_helper.py b/test/onnx/test_pytorch_helper.py
@@ -4,6 +4,7 @@
 import unittest
 
 import numpy as np
+import pytorch_test_common
 
 import torch.nn.init as init
 import torch.onnx
@@ -15,7 +16,7 @@
 from torch.testing._internal.common_utils import skipIfNoLapack
 
 
-class TestCaffe2Backend(common_utils.TestCase):
+class TestCaffe2Backend(pytorch_test_common.ExportTestCase):
     @skipIfNoLapack
     @unittest.skip("test broken because Lapack was always missing.")
     def test_helper(self):
diff --git a/test/onnx/test_pytorch_jit_onnx.py b/test/onnx/test_pytorch_jit_onnx.py
@@ -1,5 +1,6 @@
 # Owner(s): ["module: onnx"]
 import onnxruntime
+import pytorch_test_common
 
 import torch
 from pytorch_test_common import skipIfNoCuda
@@ -171,7 +172,7 @@ def MakeTestCase(opset_version: int) -> type:
     name = f"TestJITIRToONNX_opset{opset_version}"
     return type(
         str(name),
-        (common_utils.TestCase,),
+        (pytorch_test_common.ExportTestCase,),
         dict(_TestJITIRToONNX.__dict__, opset_version=opset_version),
     )
 
diff --git a/test/onnx/test_pytorch_onnx_caffe2.py b/test/onnx/test_pytorch_onnx_caffe2.py
@@ -12,6 +12,7 @@
 import model_defs.word_language_model as word_language_model
 import numpy as np
 import onnx
+import pytorch_test_common
 import torch.onnx
 import torch.onnx.operators
 import torch.utils.model_zoo as model_zoo
@@ -129,18 +130,10 @@ def do_export(model, inputs, *args, **kwargs):
 }
 
 
-class TestCaffe2Backend_opset9(common_utils.TestCase):
+class TestCaffe2Backend_opset9(pytorch_test_common.ExportTestCase):
     opset_version = 9
     embed_params = False
 
-    def setUp(self):
-        # the following should ideally be super().setUp(), https://github.com/pytorch/pytorch/issues/79630
-        common_utils.TestCase.setUp(self)
-        torch.manual_seed(0)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(0)
-        np.random.seed(seed=0)
-
     def convert_cuda(self, model, input):
         cuda_model = model.cuda()
         # input might be nested - we want to move everything to GPU
@@ -3198,52 +3191,52 @@ def setup_rnn_tests():
 # to embed_params=True
 TestCaffe2BackendEmbed_opset9 = type(
     "TestCaffe2BackendEmbed_opset9",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(TestCaffe2Backend_opset9.__dict__, embed_params=True),
 )
 
 # opset 7 tests
 TestCaffe2Backend_opset7 = type(
     "TestCaffe2Backend_opset7",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(TestCaffe2Backend_opset9.__dict__, opset_version=7),
 )
 TestCaffe2BackendEmbed_opset7 = type(
     "TestCaffe2BackendEmbed_opset7",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(TestCaffe2Backend_opset9.__dict__, embed_params=True, opset_version=7),
 )
 
 # opset 8 tests
 TestCaffe2Backend_opset8 = type(
     "TestCaffe2Backend_opset8",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(TestCaffe2Backend_opset9.__dict__, opset_version=8),
 )
 TestCaffe2BackendEmbed_opset8 = type(
     "TestCaffe2BackendEmbed_opset8",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(TestCaffe2Backend_opset9.__dict__, embed_params=True, opset_version=8),
 )
 
 # opset 10 tests
 TestCaffe2Backend_opset10 = type(
     "TestCaffe2Backend_opset10",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(TestCaffe2Backend_opset9.__dict__, opset_version=10),
 )
 
 TestCaffe2BackendEmbed_opset10 = type(
     "TestCaffe2BackendEmbed_opset10",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(TestCaffe2Backend_opset9.__dict__, embed_params=True, opset_version=10),
 )
 
 # add the same test suite as above, but switch embed_params=False
 # to embed_params=True
 TestCaffe2BackendEmbed_opset9_new_jit_API = type(
     "TestCaffe2BackendEmbed_opset9_new_jit_API",
-    (common_utils.TestCase,),
+    (pytorch_test_common.ExportTestCase,),
     dict(TestCaffe2Backend_opset9.__dict__, embed_params=True),
 )
 
diff --git a/test/onnx/test_pytorch_onnx_caffe2_quantized.py b/test/onnx/test_pytorch_onnx_caffe2_quantized.py
@@ -6,13 +6,14 @@
 
 import numpy as np
 import onnx
+import pytorch_test_common
 import torch.ao.nn.quantized as nnq
 import torch.nn as nn
 import torch.onnx
 from torch.testing._internal import common_utils
 
 
-class TestQuantizedOps(common_utils.TestCase):
+class TestQuantizedOps(pytorch_test_common.ExportTestCase):
     def generic_test(
         self, model, sample_inputs, input_names=None, decimal=3, relaxed_check=False
     ):
diff --git a/test/onnx/test_pytorch_onnx_no_runtime.py b/test/onnx/test_pytorch_onnx_no_runtime.py
@@ -11,9 +11,9 @@
 from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
-
 import onnx
 import onnx.numpy_helper
+import pytorch_test_common
 
 import torch
 import torch.nn.functional as F
@@ -74,7 +74,7 @@ def export_to_onnx(
     return onnx_model
 
 
-class TestONNXExport(common_utils.TestCase):
+class TestONNXExport(pytorch_test_common.ExportTestCase):
     def test_fuse_addmm(self):
         class AddmmModel(torch.nn.Module):
             def forward(self, x):
diff --git a/test/onnx/test_pytorch_onnx_shape_inference.py b/test/onnx/test_pytorch_onnx_shape_inference.py
diff --git a/test/onnx/test_utility_funs.py b/test/onnx/test_utility_funs.py
