diff --git a/model_zoo/official/cv/ssd/README.md b/model_zoo/official/cv/ssd/README.md index 1105b42fc15..f8929630e29 100644 --- a/model_zoo/official/cv/ssd/README.md +++ b/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) diff --git a/model_zoo/official/cv/ssd/export.py b/model_zoo/official/cv/ssd/export.py index 3e2dd293885..0ef9531f47c 100644 --- a/model_zoo/official/cv/ssd/export.py +++ b/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) diff --git a/model_zoo/official/cv/ssd/src/config_ssd300.py b/model_zoo/official/cv/ssd/src/config_ssd300.py index 8c46e86c5cf..f7b921150d3 100644 --- a/model_zoo/official/cv/ssd/src/config_ssd300.py +++ b/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": "" }) diff --git a/model_zoo/official/cv/ssd/src/config_ssd_mobilenet_v1_fpn.py b/model_zoo/official/cv/ssd/src/config_ssd_mobilenet_v1_fpn.py index 8904737ad61..8eb2452b8a4 100644 --- a/model_zoo/official/cv/ssd/src/config_ssd_mobilenet_v1_fpn.py +++ b/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": "" }) diff --git a/model_zoo/official/cv/yolov4/export.py b/model_zoo/official/cv/yolov4/export.py index f4204260009..afe2d258e26 100644 --- a/model_zoo/official/cv/yolov4/export.py +++ b/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) - -def save_air(): - """Save mindir file""" - print('============= YOLOV4 start save air ==================') - - 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') - - 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) - - tensor_input_data = Tensor(input_data) - export(network, tensor_input_data, input_shape, file_name='yolov4.air', file_format='AIR') - - print("export model success.") +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() +context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args.device_id) if __name__ == "__main__": - save_air() + ts_shape = args.testing_shape + + network = YOLOV4CspDarkNet53(is_training=False) + + param_dict = load_checkpoint(args.ckpt_file) + load_param_into_net(network, param_dict) + + 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) + + export(network, input_data, input_shape, file_name=args.file_name, file_format=args.file_format)