diff --git a/tests/ut/python/onnx/__init__.py b/tests/ut/python/onnx/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/ut/python/onnx/test_onnx.py b/tests/ut/python/onnx/test_onnx.py new file mode 100644 index 00000000000..9586ec86c6b --- /dev/null +++ b/tests/ut/python/onnx/test_onnx.py @@ -0,0 +1,210 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""ut for model serialize(save/load)""" +import os +import stat +import time + +import numpy as np +import pytest + +import mindspore.common.dtype as mstype +import mindspore.nn as nn +from mindspore import context +from mindspore.common.parameter import Parameter +from mindspore.common.tensor import Tensor +from mindspore.ops import operations as P +from mindspore.train.serialization import export + +context.set_context(mode=context.GRAPH_MODE) + +def is_enable_onnxruntime(): + val = os.getenv("ENABLE_ONNXRUNTIME", "False") + if val in ('ON', 'on', 'TRUE', 'True', 'true'): + return True + return False + +run_on_onnxruntime = pytest.mark.skipif(not is_enable_onnxruntime(), reason="Only support running on onnxruntime") + + +def setup_module(): + pass + + +def teardown_module(): + cur_dir = os.path.dirname(os.path.realpath(__file__)) + for filename in os.listdir(cur_dir): + if filename.find('ms_output_') == 0 and filename.find('.pb') > 0: + # delete temp files generated by run ut + os.chmod(filename, stat.S_IWRITE) + os.remove(filename) + + +class BatchNormTester(nn.Cell): + "used to test exporting network in training mode in onnx format" + + def __init__(self, num_features): + super(BatchNormTester, self).__init__() + self.bn = nn.BatchNorm2d(num_features) + + def construct(self, x): + return self.bn(x) + + +def test_batchnorm_train_onnx_export(): + "test onnx export interface does not modify trainable flag of a network" + input = Tensor(np.ones([1, 3, 32, 32]).astype(np.float32) * 0.01) + net = BatchNormTester(3) + net.set_train() + if not net.training: + raise ValueError('netowrk is not in training mode') + onnx_file = 'batch_norm.onnx' + export(net, input, file_name=onnx_file, file_format='ONNX') + + if not net.training: + raise ValueError('netowrk is not in training mode') + # check existence of exported onnx file and delete it + assert os.path.exists(onnx_file) + os.chmod(onnx_file, stat.S_IWRITE) + os.remove(onnx_file) + +class LeNet5(nn.Cell): + """LeNet5 definition""" + + def __init__(self): + super(LeNet5, self).__init__() + self.conv1 = nn.Conv2d(1, 6, 5, pad_mode='valid') + self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid') + self.fc1 = nn.Dense(16 * 5 * 5, 120) + self.fc2 = nn.Dense(120, 84) + self.fc3 = nn.Dense(84, 10) + self.relu = nn.ReLU() + self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2) + self.flatten = P.Flatten() + + def construct(self, x): + x = self.max_pool2d(self.relu(self.conv1(x))) + x = self.max_pool2d(self.relu(self.conv2(x))) + x = self.flatten(x) + x = self.relu(self.fc1(x)) + x = self.relu(self.fc2(x)) + x = self.fc3(x) + return x + + +class DefinedNet(nn.Cell): + """simple Net definition with maxpoolwithargmax.""" + + def __init__(self, num_classes=10): + super(DefinedNet, self).__init__() + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=0, weight_init="zeros") + self.bn1 = nn.BatchNorm2d(64) + self.relu = nn.ReLU() + self.maxpool = P.MaxPoolWithArgmax(padding="same", ksize=2, strides=2) + self.flatten = nn.Flatten() + self.fc = nn.Dense(int(56 * 56 * 64), num_classes) + + def construct(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x, argmax = self.maxpool(x) + x = self.flatten(x) + x = self.fc(x) + return x + + +class DepthwiseConv2dAndReLU6(nn.Cell): + "Net for testing DepthwiseConv2d and ReLU6" + + def __init__(self, input_channel, kernel_size): + super(DepthwiseConv2dAndReLU6, self).__init__() + weight_shape = [1, input_channel, kernel_size, kernel_size] + from mindspore.common.initializer import initializer + self.weight = Parameter(initializer('ones', weight_shape), name='weight') + self.depthwise_conv = P.DepthwiseConv2dNative(channel_multiplier=1, kernel_size=(kernel_size, kernel_size)) + self.relu6 = nn.ReLU6() + + def construct(self, x): + x = self.depthwise_conv(x, self.weight) + x = self.relu6(x) + return x + + +# generate mindspore Tensor by shape and numpy datatype +def gen_tensor(shape, dtype=np.float32): + return Tensor(np.ones(shape).astype(dtype)) + + +# ut configs in triple: (ut_name, network, network-input) +net_cfgs = [ + ('lenet', LeNet5(), gen_tensor([1, 1, 32, 32])), + ('maxpoolwithargmax', DefinedNet(), gen_tensor([1, 3, 224, 224])), + ('depthwiseconv_relu6', DepthwiseConv2dAndReLU6(3, kernel_size=3), gen_tensor([1, 3, 32, 32])), +] + + +def get_id(cfg): + return list(map(lambda x: x[0], net_cfgs)) + + +# use `pytest test_onnx.py::test_onnx_export[name]` or `pytest test_onnx.py::test_onnx_export -k name` to run single ut +@pytest.mark.parametrize('name, net, inp', net_cfgs, ids=get_id(net_cfgs)) +def test_onnx_export(name, net, inp): + onnx_file = name + ".onnx" + export(net, inp, file_name=onnx_file, file_format='ONNX') + + # check existence of exported onnx file and delete it + assert os.path.exists(onnx_file) + os.chmod(onnx_file, stat.S_IWRITE) + os.remove(onnx_file) + + +@run_on_onnxruntime +@pytest.mark.parametrize('name, net, inp', net_cfgs, ids=get_id(net_cfgs)) +def test_onnx_export_load_run(name, net, inp): + onnx_file = name + ".onnx" + export(net, inp, file_name=onnx_file, file_format='ONNX') + + import onnx + import onnxruntime as ort + + print('--------------------- onnx load ---------------------') + # Load the ONNX model + model = onnx.load(onnx_file) + # Check that the IR is well formed + onnx.checker.check_model(model) + # Print a human readable representation of the graph + g = onnx.helper.printable_graph(model.graph) + print(g) + + print('------------------ onnxruntime run ------------------') + ort_session = ort.InferenceSession(onnx_file) + input_map = {'x': inp.asnumpy()} + # provide only input x to run model + outputs = ort_session.run(None, input_map) + print(outputs[0]) + # overwrite default weight to run model + for item in net.trainable_params(): + default_value = item.default_input.asnumpy() + input_map[item.name] = np.ones(default_value.shape, dtype=default_value.dtype) + outputs = ort_session.run(None, input_map) + print(outputs[0]) + + # check existence of exported onnx file and delete it + assert os.path.exists(onnx_file) + os.chmod(onnx_file, stat.S_IWRITE) + os.remove(onnx_file) + diff --git a/tests/ut/python/utils/test_serialize.py b/tests/ut/python/utils/test_serialize.py index 616a70bfa7e..637161582a6 100644 --- a/tests/ut/python/utils/test_serialize.py +++ b/tests/ut/python/utils/test_serialize.py @@ -321,171 +321,8 @@ def test_export(): export(net, input_data, file_name="./me_export.pb", file_format="GEIR") -class BatchNormTester(nn.Cell): - "used to test exporting network in training mode in onnx format" - - def __init__(self, num_features): - super(BatchNormTester, self).__init__() - self.bn = nn.BatchNorm2d(num_features) - - def construct(self, x): - return self.bn(x) - - -class DepthwiseConv2dAndReLU6(nn.Cell): - "Net for testing DepthwiseConv2d and ReLU6" - - def __init__(self, input_channel, kernel_size): - super(DepthwiseConv2dAndReLU6, self).__init__() - weight_shape = [1, input_channel, kernel_size, kernel_size] - from mindspore.common.initializer import initializer - self.weight = Parameter(initializer('ones', weight_shape), name='weight') - self.depthwise_conv = P.DepthwiseConv2dNative(channel_multiplier=1, kernel_size=(kernel_size, kernel_size)) - self.relu6 = nn.ReLU6() - - def construct(self, x): - x = self.depthwise_conv(x, self.weight) - x = self.relu6(x) - return x - - -def test_batchnorm_train_onnx_export(): - input = Tensor(np.ones([1, 3, 32, 32]).astype(np.float32) * 0.01) - net = BatchNormTester(3) - net.set_train() - if not net.training: - raise ValueError('netowrk is not in training mode') - export(net, input, file_name='batch_norm.onnx', file_format='ONNX') - if not net.training: - raise ValueError('netowrk is not in training mode') - - -class LeNet5(nn.Cell): - """LeNet5 definition""" - - def __init__(self): - super(LeNet5, self).__init__() - self.conv1 = nn.Conv2d(1, 6, 5, pad_mode='valid') - self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid') - self.fc1 = nn.Dense(16 * 5 * 5, 120) - self.fc2 = nn.Dense(120, 84) - self.fc3 = nn.Dense(84, 10) - self.relu = nn.ReLU() - self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2) - self.flatten = P.Flatten() - - def construct(self, x): - x = self.max_pool2d(self.relu(self.conv1(x))) - x = self.max_pool2d(self.relu(self.conv2(x))) - x = self.flatten(x) - x = self.relu(self.fc1(x)) - x = self.relu(self.fc2(x)) - x = self.fc3(x) - return x - - -def test_lenet5_onnx_export(): - input = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32) * 0.01) - net = LeNet5() - export(net, input, file_name='lenet5.onnx', file_format='ONNX') - - -class DefinedNet(nn.Cell): - """simple Net definition with maxpoolwithargmax.""" - - def __init__(self, num_classes=10): - super(DefinedNet, self).__init__() - self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=0, weight_init="zeros") - self.bn1 = nn.BatchNorm2d(64) - self.relu = nn.ReLU() - self.maxpool = P.MaxPoolWithArgmax(padding="same", ksize=2, strides=2) - self.flatten = nn.Flatten() - self.fc = nn.Dense(int(56 * 56 * 64), num_classes) - - def construct(self, x): - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - x, argmax = self.maxpool(x) - x = self.flatten(x) - x = self.fc(x) - return x - - -def test_net_onnx_maxpoolwithargmax_export(): - input = Tensor(np.ones([1, 3, 224, 224]).astype(np.float32) * 0.01) - net = DefinedNet() - export(net, input, file_name='definedNet.onnx', file_format='ONNX') - - -@run_on_onnxruntime -def test_lenet5_onnx_load_run(): - onnx_file = 'lenet5.onnx' - - input = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32) * 0.01) - net = LeNet5() - export(net, input, file_name=onnx_file, file_format='ONNX') - - import onnx - import onnxruntime as ort - - print('--------------------- onnx load ---------------------') - # Load the ONNX model - model = onnx.load(onnx_file) - # Check that the IR is well formed - onnx.checker.check_model(model) - # Print a human readable representation of the graph - g = onnx.helper.printable_graph(model.graph) - print(g) - - print('------------------ onnxruntime run ------------------') - ort_session = ort.InferenceSession(onnx_file) - input_map = {'x': input.asnumpy()} - # provide only input x to run model - outputs = ort_session.run(None, input_map) - print(outputs[0]) - # overwrite default weight to run model - for item in net.trainable_params(): - input_map[item.name] = np.ones(item.default_input.asnumpy().shape, dtype=np.float32) - outputs = ort_session.run(None, input_map) - print(outputs[0]) - - -@run_on_onnxruntime -def test_depthwiseconv_relu6_onnx_load_run(): - onnx_file = 'depthwiseconv_relu6.onnx' - input_channel = 3 - input = Tensor(np.ones([1, input_channel, 32, 32]).astype(np.float32) * 0.01) - net = DepthwiseConv2dAndReLU6(input_channel, kernel_size=3) - export(net, input, file_name=onnx_file, file_format='ONNX') - - import onnx - import onnxruntime as ort - - print('--------------------- onnx load ---------------------') - # Load the ONNX model - model = onnx.load(onnx_file) - # Check that the IR is well formed - onnx.checker.check_model(model) - # Print a human readable representation of the graph - g = onnx.helper.printable_graph(model.graph) - print(g) - - print('------------------ onnxruntime run ------------------') - ort_session = ort.InferenceSession(onnx_file) - input_map = {'x': input.asnumpy()} - # provide only input x to run model - outputs = ort_session.run(None, input_map) - print(outputs[0]) - # overwrite default weight to run model - for item in net.trainable_params(): - input_map[item.name] = np.ones(item.default_input.asnumpy().shape, dtype=np.float32) - outputs = ort_session.run(None, input_map) - print(outputs[0]) - - def teardown_module(): - files = ['parameters.ckpt', 'new_ckpt.ckpt', 'lenet5.onnx', 'batch_norm.onnx', 'empty.ckpt'] + files = ['parameters.ckpt', 'new_ckpt.ckpt', 'empty.ckpt'] for item in files: file_name = './' + item if not os.path.exists(file_name):