Add export.py to ssd_ghostnet

model_zoo/research/cv/ssd_ghostnet/export.py

@@ -0,0 +1,48 @@
+# 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"""
+
+import argparse
+import numpy as np
+from mindspore import Tensor
+from mindspore import context
+from mindspore.train.serialization import load_checkpoint, load_param_into_net, export
+
+from src.ssd_ghostnet import SSD300, ssd_ghostnet
+from src.config_ghostnet_13x import config
+
+parser = argparse.ArgumentParser(description="openpose export")
+parser.add_argument("--device_id", type=int, default=0, help="Device id")
+parser.add_argument("--batch_size", type=int, default=1, help="batch size")
+parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
+parser.add_argument("--file_name", type=str, default="ssd_ghostnet", 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, default="Ascend",
+                    choices=["Ascend", "GPU", "CPU"], help="device target (default: Ascend)")
+args = parser.parse_args()
+
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
+
+if __name__ == "__main__":
+    context.set_context(mode=context.GRAPH_MODE, save_graphs=False)
+    # define net
+    net = SSD300(ssd_ghostnet(), config, is_training=False)
+
+    # load checkpoint
+    param_dict = load_checkpoint(args.ckpt_file)
+    load_param_into_net(net, param_dict)
+    input_shape = config["img_shape"]
+    inputs = np.ones([args.batch_size, 3, input_shape[0], input_shape[1]]).astype(np.float32)
+    export(net, Tensor(inputs), file_name=args.file_name, file_format=args.file_format)