Improve onnx test suite (#19662)

Author: Zha0q1

ci/docker/install/ubuntu_onnx.sh

@@ -31,4 +31,4 @@ apt-get update || true
 apt-get install -y libprotobuf-dev protobuf-compiler
 
 echo "Installing pytest, pytest-cov, protobuf, Pillow, ONNX, tabulate and onnxruntime..."
-pip3 install pytest==3.6.3 pytest-cov==2.5.1 protobuf==3.5.2 onnx==1.7.0 Pillow==5.0.0 tabulate==0.7.5 onnxruntime==1.4.0
+pip3 install pytest pytest-cov protobuf==3.5.2 onnx==1.7.0 Pillow==5.0.0 tabulate==0.7.5 onnxruntime==1.4.0

tests/python-pytest/onnx/test_operators.py

@@ -23,17 +23,32 @@
 import pytest
 import tempfile
 
-def op_export_test(op_name, Model, inputs, tmp_path):
-    def export_to_onnx(model, op_name, inputs):
-        model_path = '{}/{}'.format(tmp_path, op_name)
+def def_model(op_name, **params):
+    class Model(HybridBlock):
+        def __init__(self, **kwargs):
+            super(Model, self).__init__(**kwargs)
+
+        def hybrid_forward(self, F, *inputs):
+            names = op_name.split('.')
+            func = F
+            for name in names:
+                func = getattr(func, name)
+            out = func(*inputs, **params)
+            return out
+    return Model
+
+def op_export_test(model_name, Model, inputs, tmp_path):
+    def export_to_onnx(model, model_name, inputs):
+        model_path = '{}/{}'.format(tmp_path, model_name)
         model.export(model_path, epoch=0)
         sym_file = '{}-symbol.json'.format(model_path)
         params_file = '{}-0000.params'.format(model_path)
         dtype = inputs[0].dtype
-        onnx_file = '{}/{}.onnx'.format(tmp_path, op_name)
-        mx.contrib.onnx.export_model(sym_file, params_file, [i.shape for i in inputs],
+        onnx_file = '{}/{}.onnx'.format(tmp_path, model_name)
+        mx.contrib.onnx.export_model(sym_file, params_file, [inp.shape for inp in inputs],
                                      dtype, onnx_file)
         return onnx_file
+
     def onnx_rt(onnx_file, inputs):
         sess = rt.InferenceSession(onnx_file)
         input_dict = dict((sess.get_inputs()[i].name, inputs[i].asnumpy()) for i in range(len(inputs)))
@@ -45,90 +60,47 @@ def onnx_rt(onnx_file, inputs):
     model.initialize(ctx=mx.cpu(0))
     model.hybridize()
     pred_nat = model(*inputs)
-    onnx_file = export_to_onnx(model, op_name, inputs)
+    onnx_file = export_to_onnx(model, model_name, inputs)
     pred_onx = onnx_rt(onnx_file, inputs)
     assert_almost_equal(pred_nat, pred_onx)
 
 
-def test_onnx_export_abs():
-    with tempfile.TemporaryDirectory() as tmp_path:
-        class Model(HybridBlock):
-            def __init__(self, **kwargs):
-                super(Model, self).__init__(**kwargs)
-            def hybrid_forward(self, F, x):
-                out = F.abs(x)
-                return out
-        x = mx.nd.array([[-2, -1], [0, 99]], dtype='float32')
-        op_export_test('abs', Model, [x], tmp_path)
-
-def test_onnx_export_slice():
-    with tempfile.TemporaryDirectory() as tmp_path:
-        class Model(HybridBlock):
-            def __init__(self, **kwargs):
-                super(Model, self).__init__(**kwargs)
-            def hybrid_forward(self, F, x):
-                out = F.slice(x, begin=(0,1), end=(2,4))
-                return out
-        x = mx.nd.array([[1,2,3,4],[5,6,7,8],[9,10,11,12]], dtype='float32')
-        op_export_test('slice', Model, [x], tmp_path)
-
-def test_onnx_export_stack():
-    with tempfile.TemporaryDirectory() as tmp_path:
-        dtype = 'float32'
-        class Model(HybridBlock):
-            def __init__(self, **kwargs):
-                super(Model, self).__init__(**kwargs)
-            def hybrid_forward(self, F, x, y):
-                out = F.stack(x, y)
-                return out
-        x = mx.nd.array([1, 2], dtype=dtype)
-        y = mx.nd.array([3, 4], dtype=dtype)
-        op_export_test('stack', Model, [x, y], tmp_path)
-
-def test_onnx_export_zeros_like():
-    with tempfile.TemporaryDirectory() as tmp_path:
-        class Model(HybridBlock):
-            def __init__(self, **kwargs):
-                super(Model, self).__init__(**kwargs)
-            def hybrid_forward(self, F, x):
-                out = F.zeros_like(x)
-                return out
-        x = mx.nd.array([[-2,-1,0],[0,50,99],[4,5,6],[7,8,9]], dtype='float32')
-        op_export_test('zeros_like', Model, [x], tmp_path)
+def test_onnx_export_abs(tmp_path):
+    M = def_model('abs')
+    x = mx.nd.array([[-2, -1], [0, 99]], dtype='float32')
+    op_export_test('abs', M, [x], tmp_path)
+
+
+def test_onnx_export_slice(tmp_path):
+    M = def_model('slice', begin=(0,1), end=(2,4))
+    x = mx.nd.array([[1,2,3,4],[5,6,7,8],[9,10,11,12]], dtype='float32')
+    op_export_test('slice', M, [x], tmp_path)
+
+
+def test_onnx_export_stack(tmp_path):
+    M = def_model('stack')
+    x = mx.nd.array([1, 2], dtype='float32')
+    y = mx.nd.array([3, 4], dtype='float32')
+    op_export_test('stack', M, [x, y], tmp_path)
+
+
+def test_onnx_export_zeros_like(tmp_path):
+    M = def_model('zeros_like')
+    x = mx.nd.array([[-2,-1,0],[0,50,99],[4,5,6],[7,8,9]], dtype='float32')
+    op_export_test('zeros_like', M, [x], tmp_path)
+
 
 @pytest.mark.parametrize("dtype", ["float32", "double"])
-def test_onnx_export_arange_like(dtype):
-    with tempfile.TemporaryDirectory() as tmp_path:
-        class Model(HybridBlock):
-            def __init__(self, **kwargs):
-                super(Model, self).__init__(**kwargs)
-            def hybrid_forward(self, F, x):
-                out = F.contrib.arange_like(x)
-                return out
-        x = mx.nd.array([[-2,-1,0],[0,50,99],[4,5,6],[7,8,9]], dtype=dtype)
-        op_export_test('arange_like', Model, [x], tmp_path)
-
-def test_onnx_export_layernorm():
-    with tempfile.TemporaryDirectory() as tmp_path:
-        dtype = 'float32'
-        class Model(HybridBlock):
-            def __init__(self, **kwargs):
-                super(Model, self).__init__(**kwargs)
-            def hybrid_forward(self, F, x, gamma, beta):
-                out = F.LayerNorm(x, gamma, beta, axis=1)
-                return out
-        x = mx.nd.array([[1,3],[2,4]], dtype=dtype)
-        gamma = mx.random.uniform(0, 1, x[0].shape).astype(dtype)
-        beta = mx.random.uniform(0, 1, x[0].shape).astype(dtype)
-        op_export_test('LayerNorm', Model, [x, gamma, beta], tmp_path)
-
-
-if __name__ == '__main__':
-    test_onnx_export_abs()
-    test_onnx_export_slice()
-    test_onnx_export_stack()
-    test_onnx_export_zeros_like()
-    test_onnx_export_arange_like('float32')
-    test_onnx_export_arange_like('double')
-    test_onnx_export_layernorm()
+def test_onnx_export_arange_like(tmp_path, dtype):
+    M = def_model('contrib.arange_like')
+    x = mx.nd.array([[-2,-1,0],[0,50,99],[4,5,6],[7,8,9]], dtype=dtype)
+    op_export_test('arange_like', M, [x], tmp_path)
+
+
+def test_onnx_export_layernorm(tmp_path):
+    M = def_model('LayerNorm', axis=1)
+    x = mx.nd.array([[1,3],[2,4]], dtype='float32')
+    gamma = mx.random.uniform(0, 1, x[0].shape, dtype='float32')
+    beta = mx.random.uniform(0, 1, x[0].shape, dtype='float32')
+    op_export_test('LayerNorm', M, [x, gamma, beta], tmp_path)