!8743 modify export script for ssd and yolov4

From: @yuzhenhua666
Reviewed-by: @c_34
Signed-off-by: @c_34

model_zoo/official/cv/ssd/README.md

@@ -333,11 +333,10 @@ mAP: 0.2244936111705981
 
 ## [Export MindIR](#contents)
 
-Change the export mode and export file in `src/config.py`, and run `export.py`.
-
 ```
-python export.py --run_platform [PLATFORM] --checkpoint_path [CKPT_PATH]
+python export.py --ckpt_file [CKPT_PATH] --file_name [FILE_NAME] --file_format [FILE_FORMAT]
 ```
+The ckpt_file parameter is required.
 
 # [Model Description](#contents)
 ## [Performance](#contents)

model_zoo/official/cv/ssd/export.py

@@ -12,29 +12,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""
-ssd export mindir.
-"""
+
 import argparse
 import numpy as np
-from mindspore import context, Tensor, load_checkpoint, load_param_into_net, export
-from src.ssd import SSD300, ssd_mobilenet_v2
+
+import mindspore
+from mindspore import context, Tensor
+from mindspore.train.serialization import load_checkpoint, load_param_into_net, export
+from src.ssd import SSD300, ssd_mobilenet_v2, ssd_mobilenet_v1_fpn
 from src.config import config
 
-def get_export_args():
-    parser = argparse.ArgumentParser(description='SSD export')
-    parser.add_argument("--checkpoint_path", type=str, required=True, help="Checkpoint file path.")
-    parser.add_argument("--run_platform", type=str, default="Ascend", choices=("Ascend", "GPU", "CPU"),
-                        help="run platform, support Ascend, GPU and CPU.")
-    return parser.parse_args()
+parser = argparse.ArgumentParser(description='SSD 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.air", help="output file name.")
+parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='AIR', help='file format')
+args = parser.parse_args()
+
+context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args.device_id)
 
 if __name__ == '__main__':
-    args_opt = get_export_args()
-    context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.run_platform)
-    net = SSD300(ssd_mobilenet_v2(), config, is_training=False)
+    if config.model == "ssd300":
+        net = SSD300(ssd_mobilenet_v2(), config, is_training=False)
+    else:
+        net = ssd_mobilenet_v1_fpn(config=config)
 
-    param_dict = load_checkpoint(args_opt.checkpoint_path)
+    param_dict = load_checkpoint(args.ckpt_file)
+    net.init_parameters_data()
     load_param_into_net(net, param_dict)
-    input_shp = [1, 3] + config.img_shape
-    input_array = Tensor(np.random.uniform(-1.0, 1.0, size=input_shp).astype(np.float32))
-    export(net, input_array, file_name=config.export_file, file_format=config.export_format)
+    net.set_train(False)
+
+    input_shp = [args.batch_size, 3] + config.img_shape
+    input_array = Tensor(np.random.uniform(-1.0, 1.0, size=input_shp), mindspore.float32)
+    export(net, input_array, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/ssd/src/config_ssd300.py

@@ -78,7 +78,5 @@ config = ed({
     "voc_root": "/data/voc_dataset",
     # if coco or voc used, `image_dir` and `anno_path` are useless.
     "image_dir": "",
-    "anno_path": "",
-    "export_format": "MINDIR",
-    "export_file": "ssd.mindir"
+    "anno_path": ""
 })

model_zoo/official/cv/ssd/src/config_ssd_mobilenet_v1_fpn.py

@@ -82,7 +82,5 @@ config = ed({
     "voc_root": "/data/voc_dataset",
     # if coco or voc used, `image_dir` and `anno_path` are useless.
     "image_dir": "",
-    "anno_path": "",
-    "export_format": "MINDIR",
-    "export_file": "ssd.mindir"
+    "anno_path": ""
 })

model_zoo/official/cv/yolov4/export.py

@@ -12,54 +12,35 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""Convert ckpt to air."""
-import os
 import argparse
-
 import numpy as np
 
 import mindspore
-from mindspore import context
-from mindspore import Tensor
+from mindspore import context, Tensor
 from mindspore.train.serialization import export, load_checkpoint, load_param_into_net
 
 from src.yolo import YOLOV4CspDarkNet53
 
-context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", save_graphs=False)
+parser = argparse.ArgumentParser(description='yolov4 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("--testing_shape", type=int, default=608, help="test shape")
+parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
+parser.add_argument("--file_name", type=str, default="ssd.air", help="output file name.")
+parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='AIR', help='file format')
+args = parser.parse_args()
 
-def save_air():
-    """Save mindir file"""
-    print('============= YOLOV4 start save air ==================')
+context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args.device_id)
 
-    parser = argparse.ArgumentParser(description='Convert ckpt to air')
-    parser.add_argument('--pretrained', type=str, default='', help='pretrained model to load')
-    parser.add_argument('--batch_size', type=int, default=8, help='batch size')
+if __name__ == "__main__":
+    ts_shape = args.testing_shape
 
-    args = parser.parse_args()
     network = YOLOV4CspDarkNet53(is_training=False)
-    input_shape = Tensor(tuple([416, 416]), mindspore.float32)
-    if os.path.isfile(args.pretrained):
-        param_dict = load_checkpoint(args.pretrained)
-        param_dict_new = {}
-        for key, values in param_dict.items():
-            if key.startswith('moments.'):
-                continue
-            elif key.startswith('yolo_network.'):
-                param_dict_new[key[13:]] = values
-
-            else:
-                param_dict_new[key] = values
-
-        load_param_into_net(network, param_dict_new)
-        print('load model {} success'.format(args.pretrained))
 
-        input_data = np.random.uniform(low=0, high=1.0, size=(args.batch_size, 3, 416, 416)).astype(np.float32)
+    param_dict = load_checkpoint(args.ckpt_file)
+    load_param_into_net(network, param_dict)
 
-        tensor_input_data = Tensor(input_data)
-        export(network, tensor_input_data, input_shape, file_name='yolov4.air', file_format='AIR')
+    input_shape = Tensor(tuple([ts_shape, ts_shape]), mindspore.float32)
+    input_data = Tensor(np.zeros([args.batch_size, 3, ts_shape, ts_shape]), mindspore.float32)
 
-        print("export model success.")
-
-
-if __name__ == "__main__":
-    save_air()
+    export(network, input_data, input_shape, file_name=args.file_name, file_format=args.file_format)