From 13bb96b3b193c207980c02a1d4311d64c7c67b38 Mon Sep 17 00:00:00 2001
From: Yanjun Peng <pengyanjun1@huawei.com>
Date: Thu, 24 Dec 2020 21:29:04 +0800
Subject: [PATCH] add export for shufflenetv1

---
 .../official/cv/shufflenetv1/README_CN.md     |  3 +-
 model_zoo/official/cv/shufflenetv1/export.py  | 51 +++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 model_zoo/official/cv/shufflenetv1/export.py

diff --git a/model_zoo/official/cv/shufflenetv1/README_CN.md b/model_zoo/official/cv/shufflenetv1/README_CN.md
index 27aa3e8e7e..2b40ddce0d 100644
--- a/model_zoo/official/cv/shufflenetv1/README_CN.md
+++ b/model_zoo/official/cv/shufflenetv1/README_CN.md
@@ -67,6 +67,7 @@ ShuffleNetV1的核心部分被分成三个阶段，每个阶段重复堆积了
     ├─crossentropysmooth.py                   # 损失函数定义
     ├─lr_generator.py                         # 学习率生成函数
   ├─train.py                                  # 网络训练脚本
+  ├─export.py                                 # 模型格式转换脚本
   └─eval.py                                   # 网络评估脚本
 ```
 
@@ -174,4 +175,4 @@ result:{'Loss': 2.0479587888106323, 'Top_1_Acc': 0.7385817307692307, 'Top_5_Acc'
 
 # ModelZoo
 
-请核对官方 [主页](https://gitee.com/mindspore/mindspore/tree/master/model_zoo)。
\ No newline at end of file
+请核对官方 [主页](https://gitee.com/mindspore/mindspore/tree/master/model_zoo)。
diff --git a/model_zoo/official/cv/shufflenetv1/export.py b/model_zoo/official/cv/shufflenetv1/export.py
new file mode 100644
index 0000000000..93982b0fc2
--- /dev/null
+++ b/model_zoo/official/cv/shufflenetv1/export.py
@@ -0,0 +1,51 @@
+# 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.
+# ============================================================================
+"""
+##############export checkpoint file into air, onnx, mindir models#################
+python export.py
+"""
+import argparse
+import numpy as np
+
+import mindspore as ms
+from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
+
+from src.config import config as cfg
+from src.shufflenetv1 import ShuffleNetV1
+
+parser = argparse.ArgumentParser(description='ShuffleNetV1 export')
+parser.add_argument("--device_id", type=int, default=0, help="device id")
+parser.add_argument("--ckpt_file", type=str, required=True, help="checkpoint file path.")
+parser.add_argument("--file_name", type=str, default="shufflenetv1", help="output file name.")
+parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='AIR', help='file format')
+parser.add_argument("--device_target", type=str, choices=["Ascend", "GPU", "CPU"], default="Ascend",
+                    help="device target")
+parser.add_argument('--model_size', type=str, default='2.0x', choices=['2.0x', '1.5x', '1.0x', '0.5x'],
+                    help='shufflenetv1 model size')
+
+args = parser.parse_args()
+
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
+
+if __name__ == '__main__':
+
+    net = ShuffleNetV1(model_size=args.model_size)
+
+    param_dict = load_checkpoint(args.ckpt_file)
+    load_param_into_net(net, param_dict)
+
+    image_height, image_width = (224, 224)
+    input_arr = Tensor(np.ones([cfg.batch_size, 3, image_height, image_width]), ms.float32)
+    export(net, input_arr, file_name=args.file_name, file_format=args.file_format)