diff --git a/test/onnx/test_custom_ops.py b/test/onnx/test_custom_ops.py index 2185fe8e6d1..d1ddd954469 100644 --- a/test/onnx/test_custom_ops.py +++ b/test/onnx/test_custom_ops.py @@ -1,7 +1,5 @@ # Owner(s): ["module: onnx"] -import unittest - import numpy as np import onnx from test_pytorch_onnx_caffe2 import do_export @@ -11,9 +9,10 @@ import caffe2.python.onnx.backend as c2 import torch import torch.utils.cpp_extension from torch.onnx.symbolic_helper import _unimplemented +from test_pytorch_common import TestCase, run_tests -class TestCustomOps(unittest.TestCase): +class TestCustomOps(TestCase): def test_custom_add(self): op_source = """ #include @@ -57,7 +56,7 @@ class TestCustomOps(unittest.TestCase): np.testing.assert_array_equal(caffe2_out[0], model(x, y).cpu().numpy()) -class TestCustomAutogradFunction(unittest.TestCase): +class TestCustomAutogradFunction(TestCase): opset_version = 9 keep_initializers_as_inputs = False onnx_shape_inference = True @@ -129,7 +128,7 @@ class TestCustomAutogradFunction(unittest.TestCase): run_model_test(self, model, input_args=(x,)) -class TestExportAsContribOps(unittest.TestCase): +class TestExportAsContribOps(TestCase): opset_version = 14 keep_initializers_as_inputs = False onnx_shape_inference = True @@ -164,4 +163,4 @@ class TestExportAsContribOps(unittest.TestCase): if __name__ == "__main__": - unittest.main() + run_tests() diff --git a/test/onnx/test_models_onnxruntime.py b/test/onnx/test_models_onnxruntime.py index 9370c8e0dab..783c0981173 100644 --- a/test/onnx/test_models_onnxruntime.py +++ b/test/onnx/test_models_onnxruntime.py @@ -1,9 +1,8 @@ # Owner(s): ["module: onnx"] -import unittest - import onnxruntime # noqa: F401 from test_models import TestModels +from test_pytorch_common import TestCase, run_tests from test_pytorch_onnx_onnxruntime import run_model_test import torch @@ -24,7 +23,7 @@ def exportTest(self, model, inputs, rtol=1e-2, atol=1e-7, opset_versions=None): TestModels = type( "TestModels", - (unittest.TestCase,), + (TestCase,), dict(TestModels.__dict__, is_script_test_enabled=False, exportTest=exportTest), ) @@ -32,7 +31,7 @@ TestModels = type( # model tests for scripting with new JIT APIs and shape inference TestModels_new_jit_API = type( "TestModels_new_jit_API", - (unittest.TestCase,), + (TestCase,), dict( TestModels.__dict__, exportTest=exportTest, @@ -43,4 +42,4 @@ TestModels_new_jit_API = type( if __name__ == "__main__": - unittest.main() + run_tests() diff --git a/test/onnx/test_operators.py b/test/onnx/test_operators.py index 9315a79fb4c..28018450625 100644 --- a/test/onnx/test_operators.py +++ b/test/onnx/test_operators.py @@ -16,6 +16,7 @@ from test_pytorch_common import ( TestCase, flatten, run_tests, + skipIfCaffe2, skipIfNoLapack, ) @@ -35,7 +36,6 @@ from torch.onnx.symbolic_helper import ( _get_tensor_sizes, parse_args, ) -from torch.testing._internal.common_utils import skipIfCaffe2 """Usage: python test/onnx/test_operators.py [--no-onnx] [--produce-onnx-test-data] --no-onnx: no onnx python dependence diff --git a/test/onnx/test_pytorch_helper.py b/test/onnx/test_pytorch_helper.py index 8ab5c6f3fb0..eeb5f88f17f 100644 --- a/test/onnx/test_pytorch_helper.py +++ b/test/onnx/test_pytorch_helper.py @@ -5,7 +5,7 @@ import unittest import numpy as np from pytorch_helper import PyTorchModule -from test_pytorch_common import skipIfNoLapack +from test_pytorch_common import skipIfNoLapack, run_tests, TestCase import torch.nn.init as init import torch.onnx @@ -14,7 +14,7 @@ from caffe2.python.model_helper import ModelHelper from torch import nn -class TestCaffe2Backend(unittest.TestCase): +class TestCaffe2Backend(TestCase): @skipIfNoLapack @unittest.skip("test broken because Lapack was always missing.") def test_helper(self): @@ -67,4 +67,4 @@ class TestCaffe2Backend(unittest.TestCase): if __name__ == "__main__": - unittest.main() + run_tests() diff --git a/test/onnx/test_pytorch_jit_onnx.py b/test/onnx/test_pytorch_jit_onnx.py index ace8127062e..f0cabae76b5 100644 --- a/test/onnx/test_pytorch_jit_onnx.py +++ b/test/onnx/test_pytorch_jit_onnx.py @@ -1,11 +1,10 @@ # Owner(s): ["module: onnx"] -import unittest - import onnxruntime import torch from torch._C import parse_ir from torch.onnx import verification +from test_pytorch_common import TestCase, run_tests def _jit_graph_to_onnx_model(graph, operator_export_type, opset_version): @@ -85,7 +84,7 @@ def MakeTestCase(opset_version: int) -> type: name = f"TestJITIRToONNX_opset{opset_version}" return type( str(name), - (unittest.TestCase,), + (TestCase,), dict(_TestJITIRToONNX.__dict__, opset_version=opset_version), ) @@ -93,4 +92,4 @@ def MakeTestCase(opset_version: int) -> type: TestJITIRToONNX_opset14 = MakeTestCase(14) if __name__ == "__main__": - unittest.main() + run_tests() diff --git a/test/onnx/test_pytorch_onnx_caffe2.py b/test/onnx/test_pytorch_onnx_caffe2.py index 22cc9f46ceb..32382b9d005 100644 --- a/test/onnx/test_pytorch_onnx_caffe2.py +++ b/test/onnx/test_pytorch_onnx_caffe2.py @@ -29,6 +29,8 @@ from test_pytorch_common import ( skipIfTravis, skipIfUnsupportedMinOpsetVersion, skipIfUnsupportedOpsetVersion, + TestCase, + run_tests, ) # Import various models for testing @@ -128,7 +130,7 @@ model_urls = { } -class TestCaffe2Backend_opset9(unittest.TestCase): +class TestCaffe2Backend_opset9(TestCase): opset_version = 9 embed_params = False @@ -3195,44 +3197,44 @@ setup_rnn_tests() # to embed_params=True TestCaffe2BackendEmbed_opset9 = type( "TestCaffe2BackendEmbed_opset9", - (unittest.TestCase,), + (TestCase,), dict(TestCaffe2Backend_opset9.__dict__, embed_params=True), ) # opset 7 tests TestCaffe2Backend_opset7 = type( "TestCaffe2Backend_opset7", - (unittest.TestCase,), + (TestCase,), dict(TestCaffe2Backend_opset9.__dict__, opset_version=7), ) TestCaffe2BackendEmbed_opset7 = type( "TestCaffe2BackendEmbed_opset7", - (unittest.TestCase,), + (TestCase,), dict(TestCaffe2Backend_opset9.__dict__, embed_params=True, opset_version=7), ) # opset 8 tests TestCaffe2Backend_opset8 = type( "TestCaffe2Backend_opset8", - (unittest.TestCase,), + (TestCase,), dict(TestCaffe2Backend_opset9.__dict__, opset_version=8), ) TestCaffe2BackendEmbed_opset8 = type( "TestCaffe2BackendEmbed_opset8", - (unittest.TestCase,), + (TestCase,), dict(TestCaffe2Backend_opset9.__dict__, embed_params=True, opset_version=8), ) # opset 10 tests TestCaffe2Backend_opset10 = type( "TestCaffe2Backend_opset10", - (unittest.TestCase,), + (TestCase,), dict(TestCaffe2Backend_opset9.__dict__, opset_version=10), ) TestCaffe2BackendEmbed_opset10 = type( "TestCaffe2BackendEmbed_opset10", - (unittest.TestCase,), + (TestCase,), dict(TestCaffe2Backend_opset9.__dict__, embed_params=True, opset_version=10), ) @@ -3240,9 +3242,9 @@ TestCaffe2BackendEmbed_opset10 = type( # to embed_params=True TestCaffe2BackendEmbed_opset9_new_jit_API = type( "TestCaffe2BackendEmbed_opset9_new_jit_API", - (unittest.TestCase,), + (TestCase,), dict(TestCaffe2Backend_opset9.__dict__, embed_params=True), ) if __name__ == "__main__": - unittest.main() + run_tests() diff --git a/test/onnx/test_pytorch_onnx_caffe2_quantized.py b/test/onnx/test_pytorch_onnx_caffe2_quantized.py index de8f8e13e90..5aacba78df4 100644 --- a/test/onnx/test_pytorch_onnx_caffe2_quantized.py +++ b/test/onnx/test_pytorch_onnx_caffe2_quantized.py @@ -1,8 +1,6 @@ # Owner(s): ["module: unknown"] import io -import unittest - import numpy as np import onnx @@ -10,9 +8,10 @@ import caffe2.python.onnx.backend as c2 import torch.nn as nn import torch.nn.quantized as nnq import torch.onnx +from test_pytorch_common import TestCase, run_tests -class TestQuantizedOps(unittest.TestCase): +class TestQuantizedOps(TestCase): def generic_test( self, model, sample_inputs, input_names=None, decimal=3, relaxed_check=False ): @@ -378,4 +377,4 @@ class TestQuantizedOps(unittest.TestCase): if __name__ == "__main__": - unittest.main() + run_tests() diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py index 9f79349705f..ee3f6eb917e 100644 --- a/test/onnx/test_pytorch_onnx_onnxruntime.py +++ b/test/onnx/test_pytorch_onnx_onnxruntime.py @@ -30,11 +30,13 @@ from test_pytorch_common import ( RNN_HIDDEN_SIZE, RNN_INPUT_SIZE, RNN_SEQUENCE_LENGTH, + run_tests, skipIfNoLapack, skipIfUnsupportedMaxOpsetVersion, skipIfUnsupportedMinOpsetVersion, skipIfUnsupportedOpsetVersion, skipScriptTest, + TestCase, ) from torchvision import ops from torchvision.models.detection.faster_rcnn import FastRCNNPredictor, TwoMLPHead @@ -62,9 +64,7 @@ from torch.onnx.symbolic_helper import _unimplemented _ORT_PROVIDERS = ["CPUExecutionProvider"] -def run_model_test( - test_suite: Union[_TestONNXRuntime, unittest.TestCase], *args, **kwargs -): +def run_model_test(test_suite: Union[_TestONNXRuntime, TestCase], *args, **kwargs): kwargs["ort_providers"] = _ORT_PROVIDERS kwargs["opset_version"] = test_suite.opset_version kwargs["keep_initializers_as_inputs"] = test_suite.keep_initializers_as_inputs @@ -72,7 +72,7 @@ def run_model_test( def run_model_test_with_external_data( - test_suite: Union[_TestONNXRuntime, unittest.TestCase], *args, **kwargs + test_suite: Union[_TestONNXRuntime, TestCase], *args, **kwargs ): kwargs["use_external_data"] = True return run_model_test(test_suite, *args, **kwargs) @@ -12723,7 +12723,7 @@ def MakeTestCase(opset_version: int, keep_initializers_as_inputs: bool = True) - name += "_IRv4" return type( str(name), - (unittest.TestCase,), + (TestCase,), dict( _TestONNXRuntime.__dict__, opset_version=opset_version, @@ -12762,4 +12762,4 @@ TestONNXRuntime_opset16 = MakeTestCase(16, keep_initializers_as_inputs=False) if __name__ == "__main__": - unittest.main() + run_tests() diff --git a/test/onnx/test_pytorch_onnx_onnxruntime_cuda.py b/test/onnx/test_pytorch_onnx_onnxruntime_cuda.py index 03fec231132..b5d41bbf71f 100644 --- a/test/onnx/test_pytorch_onnx_onnxruntime_cuda.py +++ b/test/onnx/test_pytorch_onnx_onnxruntime_cuda.py @@ -8,6 +8,7 @@ from test_pytorch_common import ( skipIfNoCuda, skipIfUnsupportedMinOpsetVersion, skipScriptTest, + TestCase, ) # TODO(justinchuby): Remove reference to other unit tests. @@ -18,7 +19,7 @@ from torch.cuda.amp import autocast from torch.onnx._globals import GLOBALS -class TestONNXRuntime_cuda(unittest.TestCase): +class TestONNXRuntime_cuda(TestCase): opset_version = GLOBALS.export_onnx_opset_version keep_initializers_as_inputs = True @@ -150,4 +151,5 @@ TestONNXRuntime_cuda.setUp = TestONNXRuntime.setUp TestONNXRuntime_cuda.run_test = TestONNXRuntime.run_test if __name__ == "__main__": + # TODO: convert this to use common_utils.run_tests() unittest.main(TestONNXRuntime_cuda()) diff --git a/test/onnx/test_pytorch_onnx_shape_inference.py b/test/onnx/test_pytorch_onnx_shape_inference.py index da21513cf34..6b293d65697 100644 --- a/test/onnx/test_pytorch_onnx_shape_inference.py +++ b/test/onnx/test_pytorch_onnx_shape_inference.py @@ -1,9 +1,7 @@ # Owner(s): ["module: onnx"] -import unittest - import numpy as np -from test_pytorch_common import skipIfUnsupportedMinOpsetVersion +from test_pytorch_common import run_tests, skipIfUnsupportedMinOpsetVersion, TestCase import torch from torch.onnx import _constants @@ -21,9 +19,9 @@ def expect_tensor(scalar_type, shape=None): return verify -class TestONNXShapeInference(unittest.TestCase): +class TestONNXShapeInference(TestCase): def __init__(self, *args, **kwargs): - unittest.TestCase.__init__(self, *args, **kwargs) + TestCase.__init__(self, *args, **kwargs) self.opset_version = _constants.onnx_main_opset _set_onnx_shape_inference(True) _set_opset_version(self.opset_version) @@ -273,4 +271,4 @@ class TestONNXShapeInference(unittest.TestCase): if __name__ == "__main__": - unittest.main() + run_tests()