From e5a141fde2c25597c2d16e0771c767365e3deb10 Mon Sep 17 00:00:00 2001 From: changzherui Date: Fri, 4 Dec 2020 16:26:08 +0800 Subject: [PATCH] add export mindir example --- tests/st/export/test_train_mindir.py | 114 ++++++++++++++++++ .../st/export/{test_export.py => text_air.py} | 0 tests/st/export/text_lite_mindir.py | 49 ++++++++ 3 files changed, 163 insertions(+) create mode 100644 tests/st/export/test_train_mindir.py rename tests/st/export/{test_export.py => text_air.py} (100%) create mode 100644 tests/st/export/text_lite_mindir.py diff --git a/tests/st/export/test_train_mindir.py b/tests/st/export/test_train_mindir.py new file mode 100644 index 0000000000..70579d5ed7 --- /dev/null +++ b/tests/st/export/test_train_mindir.py @@ -0,0 +1,114 @@ +# 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. +import os +import numpy as np +import pytest + +import mindspore.nn as nn +from mindspore import context +from mindspore.common.tensor import Tensor +from mindspore.common.initializer import TruncatedNormal +from mindspore.common.parameter import ParameterTuple +from mindspore.ops import operations as P +from mindspore.ops import composite as C +from mindspore.train.serialization import export + + +def weight_variable(): + return TruncatedNormal(0.02) + + +def conv(in_channels, out_channels, kernel_size, stride=1, padding=0): + weight = weight_variable() + return nn.Conv2d(in_channels, out_channels, + kernel_size=kernel_size, stride=stride, padding=padding, + weight_init=weight, has_bias=False, pad_mode="valid") + + +def fc_with_initialize(input_channels, out_channels): + weight = weight_variable() + bias = weight_variable() + return nn.Dense(input_channels, out_channels, weight, bias) + + +class LeNet5(nn.Cell): + def __init__(self): + super(LeNet5, self).__init__() + self.batch_size = 32 + self.conv1 = conv(1, 6, 5) + self.conv2 = conv(6, 16, 5) + self.fc1 = fc_with_initialize(16 * 5 * 5, 120) + self.fc2 = fc_with_initialize(120, 84) + self.fc3 = fc_with_initialize(84, 10) + self.relu = nn.ReLU() + self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2) + self.reshape = P.Reshape() + + def construct(self, x): + x = self.conv1(x) + x = self.relu(x) + x = self.max_pool2d(x) + x = self.conv2(x) + x = self.relu(x) + x = self.max_pool2d(x) + x = self.reshape(x, (self.batch_size, -1)) + x = self.fc1(x) + x = self.relu(x) + x = self.fc2(x) + x = self.relu(x) + x = self.fc3(x) + return x + + +class WithLossCell(nn.Cell): + def __init__(self, network): + super(WithLossCell, self).__init__(auto_prefix=False) + self.loss = nn.SoftmaxCrossEntropyWithLogits() + self.network = network + + def construct(self, x, label): + predict = self.network(x) + return self.loss(predict, label) + + +class TrainOneStepCell(nn.Cell): + def __init__(self, network): + super(TrainOneStepCell, self).__init__(auto_prefix=False) + self.network = network + self.network.set_train() + self.weights = ParameterTuple(network.trainable_params()) + self.optimizer = nn.Momentum(self.weights, 0.1, 0.9) + self.hyper_map = C.HyperMap() + self.grad = C.GradOperation(get_by_list=True) + + def construct(self, x, label): + weights = self.weights + grads = self.grad(self.network, weights)(x, label) + return self.optimizer(grads) + + +@pytest.mark.level0 +@pytest.mark.platform_x86_ascend_training +@pytest.mark.platform_arm_ascend_training +@pytest.mark.env_onecard +def test_export_lenet_grad_mindir(): + context.set_context(mode=context.GRAPH_MODE, device_target="Ascend") + network = LeNet5() + network.set_train() + predict = Tensor(np.ones([32, 1, 32, 32]).astype(np.float32) * 0.01) + label = Tensor(np.zeros([32, 10]).astype(np.float32)) + net = TrainOneStepCell(WithLossCell(network)) + export(net, predict, label, file_name="lenet_grad", file_format='MINDIR') + verify_name = "lenet_grad.mindir" + assert os.path.exists(verify_name) diff --git a/tests/st/export/test_export.py b/tests/st/export/text_air.py similarity index 100% rename from tests/st/export/test_export.py rename to tests/st/export/text_air.py diff --git a/tests/st/export/text_lite_mindir.py b/tests/st/export/text_lite_mindir.py new file mode 100644 index 0000000000..a8296c8d5b --- /dev/null +++ b/tests/st/export/text_lite_mindir.py @@ -0,0 +1,49 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# 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. +# ============================================================================ +"""mobilenetv2 export mindir.""" +import os +import numpy as np +import pytest + +from mindspore import Tensor +from mindspore.train.serialization import export, load_checkpoint +from mindspore import context +from model_zoo.official.cv.mobilenetv2.src.mobilenetV2 import MobileNetV2Backbone, MobileNetV2Head, mobilenet_v2 + + +context.set_context(mode=context.GRAPH_MODE, device_target="GPU") +ckpt_path = '/home/workspace/mindspore_dataset/checkpoint/mobilenetv2/mobilenetv2_gpu.ckpt' + +@pytest.mark.level0 +@pytest.mark.platform_x86_gpu_training +@pytest.mark.env_onecard +def test_export_mobilenetv2_gpu_mindir(): + backbone_net = MobileNetV2Backbone() + head_net = MobileNetV2Head(input_channel=backbone_net.out_channels, num_classes=1000) + net = mobilenet_v2(backbone_net, head_net) + load_checkpoint(ckpt_path, net) + input_tensor = Tensor(np.ones([1, 3, 224, 224]).astype(np.float32)) + export(net, input_tensor, file_name="mobilenetv2_gpu", file_format="MINDIR") + output = net(input_tensor).asnumpy().tolist() + file_obj = open('mobilenetv2_gpu_output.txt', 'w') + file_obj.write("output 1 3 224 224\n") + for num in output[0]: + file_obj.write(str(num)) + file_obj.close() + assert os.path.exists("mobilenetv2_gpu.mindir") + assert os.path.exists("mobilenetv2_gpu_output.txt")