Skip to content
This repository was archived by the owner on Nov 17, 2023. It is now read-only.

Improve onnx test suite #19662

Merged
leezu merged 2 commits into from
Dec 14, 2020
Merged

Improve onnx test suite #19662

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1875bb5
def_model helper func
Zha0q1 Dec 12, 2020
968f9ac
update pytest
Zha0q1 Dec 12, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ci/docker/install/ubuntu_onnx.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
142 changes: 57 additions & 85 deletions tests/python-pytest/onnx/test_operators.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)))
Expand All @@ -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)