!9995 add export file for shufflenet,delete export for resnet50_quant, modify export file to support mindir

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

model_zoo/official/cv/alexnet/export.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ============================================================================
 """
-##############export checkpoint file into air and onnx models#################
+##############export checkpoint file into air, onnx, mindir models#################
 python export.py
 """
 import argparse
@@ -25,18 +25,22 @@ from mindspore import context, Tensor, load_checkpoint, load_param_into_net, exp
 from src.config import alexnet_cifar10_cfg, alexnet_imagenet_cfg
 from src.alexnet import AlexNet
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Classification')
-    parser.add_argument('--dataset_name', type=str, default='cifar10', choices=['imagenet', 'cifar10'],
-                        help='please choose dataset: imagenet or cifar10.')
-    parser.add_argument('--device_target', type=str, default="Ascend",
-                        choices=['Ascend', 'GPU'],
-                        help='device where the code will be implemented (default: Ascend)')
-    parser.add_argument('--ckpt_path', type=str, default="./ckpt", help='if is test, must provide\
-                                path where the trained ckpt file')
-    args_opt = parser.parse_args()
-    context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target)
+parser = argparse.ArgumentParser(description='Classification')
+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('--dataset_name', type=str, default='cifar10', choices=['imagenet', 'cifar10'],
+                    help='please choose dataset: imagenet or cifar10.')
+parser.add_argument('--device_target', type=str, default="Ascend",
+                    choices=['Ascend', 'GPU', 'CPU'],
+                    help='device where the code will be implemented (default: Ascend)')
+parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
+parser.add_argument("--file_name", type=str, default="alexnet", help="output file name.")
+parser.add_argument("--file_format", type=str, choices=["AIR", "ONNX", "MINDIR"], default="AIR", help="file format")
+args_opt = parser.parse_args()
+
+context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)
 
+if __name__ == '__main__':
     if args_opt.dataset_name == 'cifar10':
         cfg = alexnet_cifar10_cfg
     elif args_opt.dataset_name == 'imagenet':
@@ -46,8 +50,8 @@ if __name__ == '__main__':
 
     net = AlexNet(num_classes=cfg.num_classes)
 
-    param_dict = load_checkpoint(args_opt.ckpt_path)
+    param_dict = load_checkpoint(args_opt.ckpt_file)
     load_param_into_net(net, param_dict)
 
-    input_arr = Tensor(np.random.uniform(0.0, 1.0, size=[1, 3, cfg.image_height, cfg.image_width]), ms.float32)
-    export(net, input_arr, file_name=cfg.air_name, file_format="AIR")
+    input_arr = Tensor(np.zeros([args_opt.batch_size, 3, cfg.image_height, cfg.image_width]), ms.float32)
+    export(net, input_arr, file_name=args_opt.file_name, file_format=args_opt.file_format)

model_zoo/official/cv/centerface/export.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-
+"""export checkpoint file into air, onnx, mindir models"""
 import argparse
 import numpy as np
 
@@ -29,9 +29,11 @@ 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="centerface", 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")
 args = parser.parse_args()
 
-context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args.device_id)
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
 
 if __name__ == '__main__':
     config = ConfigCenterface()

model_zoo/official/cv/cnnctc/export.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""export checkpoint file into air models"""
+"""export checkpoint file into air, onnx, mindir models"""
 import argparse
 import numpy as np
 
@@ -22,14 +22,18 @@ import mindspore.common.dtype as mstype
 from src.config import Config_CNNCTC
 from src.cnn_ctc import CNNCTC_Model
 
-context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
-
-parser = argparse.ArgumentParser(description='CNNCTC_export')
-parser.add_argument('--ckpt_file', type=str, default='./ckpts/cnn_ctc.ckpt', help='CNN&CTC ckpt file.')
-parser.add_argument('--output_file', type=str, default='cnn_ctc', help='CNN&CTC output air name.')
+parser = argparse.ArgumentParser(description="CNNCTC_export")
+parser.add_argument("--device_id", type=int, default=0, help="Device id")
+parser.add_argument("--file_name", type=str, default="cnn_ctc", help="CNN&CTC output air 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("--ckpt_file", type=str, default="./ckpts/cnn_ctc.ckpt", help="CNN&CTC ckpt file.")
 args_opt = parser.parse_args()
 
-if __name__ == '__main__':
+context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)
+
+if __name__ == "__main__":
     cfg = Config_CNNCTC()
     ckpt_path = cfg.CKPT_PATH
 
@@ -44,4 +48,4 @@ if __name__ == '__main__':
 
     input_data = Tensor(np.zeros([bs, 3, cfg.IMG_H, cfg.IMG_W]), mstype.float32)
 
-    export(net, input_data, file_name=args_opt.output_file, file_format="AIR")
+    export(net, input_data, file_name=args_opt.file_name, file_format=args_opt.file_format)

model_zoo/official/cv/deeplabv3/export.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""export AIR file."""
+"""export checkpoint file into air, onnx, mindir models"""
 import argparse
 import numpy as np
 
@@ -20,23 +20,30 @@ from mindspore import Tensor, context, load_checkpoint, load_param_into_net, exp
 
 from src.nets import net_factory
 
-context.set_context(mode=context.GRAPH_MODE, save_graphs=False)
+parser = argparse.ArgumentParser(description='checkpoint 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("--input_size", type=int, default=513, help="batch size")
+parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
+parser.add_argument("--file_name", type=str, default="deeplabv3", 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', type=str.lower, default='deeplab_v3_s8', choices=['deeplab_v3_s16', 'deeplab_v3_s8'],
+                    help='Select model structure (Default: deeplab_v3_s8)')
+parser.add_argument('--num_classes', type=int, default=21, help='the number of classes (Default: 21)')
+args = parser.parse_args()
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='checkpoint export')
-    parser.add_argument('--checkpoint', type=str.lower, default='', help='checkpoint of deeplabv3 (Default: None)')
-    parser.add_argument('--model', type=str.lower, default='deeplab_v3_s8', choices=['deeplab_v3_s16', 'deeplab_v3_s8'],
-                        help='Select model structure (Default: deeplab_v3_s8)')
-    parser.add_argument('--num_classes', type=int, default=21, help='the number of classes (Default: 21)')
-    args = parser.parse_args()
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
 
+if __name__ == '__main__':
     if args.model == 'deeplab_v3_s16':
         network = net_factory.nets_map['deeplab_v3_s16']('eval', args.num_classes, 16, True)
     else:
         network = net_factory.nets_map['deeplab_v3_s8']('eval', args.num_classes, 8, True)
-    param_dict = load_checkpoint(args.checkpoint)
+    param_dict = load_checkpoint(args.ckpt_file)
 
     # load the parameter into net
     load_param_into_net(network, param_dict)
-    input_data = np.random.uniform(0.0, 1.0, size=[32, 3, 513, 513]).astype(np.float32)
-    export(network, Tensor(input_data), file_name=args.model, file_format='AIR')
+    input_data = Tensor(np.ones([args.batch_size, 3, args.input_size, args.input_size]).astype(np.float32))
+    export(network, input_data, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/densenet121/export.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
+"""export checkpoint file into air, onnx, mindir models"""
 import argparse
 import numpy as np
 
@@ -28,9 +29,11 @@ parser.add_argument("--batch_size", type=int, default=32, help="batch size")
 parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
 parser.add_argument("--file_name", type=str, default="densenet121", 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")
 args = parser.parse_args()
 
-context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args.device_id)
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
 
 if __name__ == "__main__":
     network = DenseNet121(config.num_classes)

model_zoo/official/cv/faster_rcnn/export.py

@@ -12,22 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""export checkpoint file into air models"""
+"""export checkpoint file into air, onnx, mindir models"""
 import argparse
 import numpy as np
 
 import mindspore as ms
-from mindspore import Tensor, load_checkpoint, load_param_into_net, export
+from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
 
 from src.FasterRcnn.faster_rcnn_r50 import Faster_Rcnn_Resnet50
 from src.config import config
 
 parser = argparse.ArgumentParser(description='fasterrcnn_export')
+parser.add_argument("--device_id", type=int, default=0, help="Device id")
+parser.add_argument("--file_name", type=str, default="faster_rcnn", 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('--ckpt_file', type=str, default='', help='fasterrcnn ckpt file.')
-parser.add_argument('--output_file', type=str, default='', help='fasterrcnn output air 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=args.device_target, device_id=args.device_id)
+
 if __name__ == '__main__':
     net = Faster_Rcnn_Resnet50(config=config)
 
@@ -40,4 +45,4 @@ if __name__ == '__main__':
     gt_label = Tensor(np.random.uniform(0.0, 1.0, size=[config.test_batch_size, config.num_gts]), ms.int32)
     gt_num = Tensor(np.random.uniform(0.0, 1.0, size=[config.test_batch_size, config.num_gts]), ms.bool_)
 
-    export(net, img, img_metas, gt_bboxes, gt_label, gt_num, file_name=args.output_file, file_format=args.file_format)
+    export(net, img, img_metas, gt_bboxes, gt_label, gt_num, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/googlenet/export.py

@@ -13,27 +13,36 @@
 # limitations under the License.
 # ============================================================================
 """
-##############export checkpoint file into air and onnx models#################
+##############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
+from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
 
 from src.config import cifar_cfg, imagenet_cfg
 from src.googlenet import GoogleNet
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Classification')
-    parser.add_argument('--dataset_name', type=str, default='cifar10', choices=['imagenet', 'cifar10'],
-                        help='dataset name.')
-    args_opt = parser.parse_args()
+parser = argparse.ArgumentParser(description='Classification')
+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="googlenet", 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('--dataset_name', type=str, default='cifar10', choices=['imagenet', 'cifar10'],
+                    help='dataset name.')
+args = parser.parse_args()
+
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
 
-    if args_opt.dataset_name == 'cifar10':
+if __name__ == '__main__':
+    if args.dataset_name == 'cifar10':
         cfg = cifar_cfg
-    elif args_opt.dataset_name == 'imagenet':
+    elif args.dataset_name == 'imagenet':
         cfg = imagenet_cfg
     else:
         raise ValueError("dataset is not support.")
@@ -41,9 +50,8 @@ if __name__ == '__main__':
     net = GoogleNet(num_classes=cfg.num_classes)
 
     assert cfg.checkpoint_path is not None, "cfg.checkpoint_path is None."
-    param_dict = load_checkpoint(cfg.checkpoint_path)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(net, param_dict)
 
-    input_arr = Tensor(np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]), ms.float32)
-    export(net, input_arr, file_name=cfg.onnx_filename, file_format="ONNX")
-    export(net, input_arr, file_name=cfg.air_filename, file_format="AIR")
+    input_arr = Tensor(np.ones([args.batch_size, 3, cfg.image_height, cfg.image_width]), ms.float32)
+    export(net, input_arr, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/inceptionv3/export.py

@@ -12,26 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""
-export checkpoint file into models
-"""
+"""export checkpoint file into air, onnx, mindir models"""
 import argparse
 import numpy as np
 
 import mindspore as ms
-from mindspore import Tensor, load_checkpoint, load_param_into_net, export
+from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
 
 from src.config import config_gpu as cfg
 from src.inception_v3 import InceptionV3
 
 parser = argparse.ArgumentParser(description='inceptionv3 export')
+parser.add_argument("--device_id", type=int, default=0, help="Device id")
 parser.add_argument('--ckpt_file', type=str, required=True, help='inceptionv3 ckpt file.')
-parser.add_argument('--output_file', type=str, default='inceptionv3', help='inceptionv3 output air name.')
+parser.add_argument('--file_name', type=str, default='inceptionv3', help='inceptionv3 output air name.')
 parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='AIR', help='file format')
 parser.add_argument('--width', type=int, default=299, help='input width')
 parser.add_argument('--height', type=int, default=299, help='input height')
+parser.add_argument("--device_target", type=str, choices=["Ascend", "GPU", "CPU"], default="Ascend",
+                    help="device target")
 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 = InceptionV3(num_classes=cfg.num_classes, is_training=False)
     param_dict = load_checkpoint(args.ckpt_file)
@@ -39,4 +42,4 @@ if __name__ == '__main__':
 
     input_arr = Tensor(np.random.uniform(0.0, 1.0, size=[cfg.batch_size, 3, args.width, args.height]), ms.float32)
 
-    export(net, input_arr, file_name=args.output_file, file_format=args.file_format)
+    export(net, input_arr, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/inceptionv4/export.py

@@ -12,35 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""
-##############export checkpoint file into air and onnx models#################
-"""
+"""export checkpoint file into air, onnx, mindir models"""
 import argparse
 import numpy as np
 
 import mindspore as ms
 from mindspore import Tensor
-from mindspore.train.serialization import load_checkpoint, load_param_into_net, export
+from mindspore.train.serialization import load_checkpoint, load_param_into_net, export, context
 
 from src.config import config_ascend as config
 from src.inceptionv4 import Inceptionv4
 
-def parse_args():
-    '''parse_args'''
-    parser = argparse.ArgumentParser(description='checkpoint export')
-    parser.add_argument('--model_name', type=str, default='inceptionV4.air', help='convert model name of inceptionv4')
-    parser.add_argument('--format', type=str, default='AIR', help='convert model name of inceptionv4')
-    parser.add_argument('--checkpoint', type=str, default='', help='checkpoint of inceptionv4')
-    _args_opt = parser.parse_args()
-    return _args_opt
+parser = argparse.ArgumentParser(description='inceptionv4 export')
+parser.add_argument("--device_id", type=int, default=0, help="Device id")
+parser.add_argument('--ckpt_file', type=str, required=True, help='inceptionv4 ckpt file.')
+parser.add_argument('--file_name', type=str, default='inceptionv4', help='inceptionv4 output air name.')
+parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='AIR', help='file format')
+parser.add_argument('--width', type=int, default=299, help='input width')
+parser.add_argument('--height', type=int, default=299, help='input height')
+parser.add_argument("--device_target", type=str, choices=["Ascend", "GPU", "CPU"], default="Ascend",
+                    help="device target")
+args = parser.parse_args()
 
-if __name__ == '__main__':
-
-    args_opt = parse_args()
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
 
+if __name__ == '__main__':
     net = Inceptionv4(classes=config.num_classes)
-    param_dict = load_checkpoint(args_opt.checkpoint)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(net, param_dict)
 
-    input_arr = Tensor(np.random.uniform(0.0, 1.0, size=[1, 3, 299, 299]), ms.float32)
-    export(net, input_arr, file_name=args_opt.model_name, file_format=args_opt.format)
+    input_arr = Tensor(np.ones([config.batch_size, 3, args.width, args.height]), ms.float32)
+    export(net, input_arr, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/lenet/export.py

@@ -12,9 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""
-export network to infer `AIR` backend.
-"""
+"""export checkpoint file into air, onnx, mindir models"""
 
 import argparse
 import numpy as np
@@ -25,25 +23,26 @@ from mindspore import Tensor, context, load_checkpoint, load_param_into_net, exp
 from src.config import mnist_cfg as cfg
 from src.lenet import LeNet5
 
-
 parser = argparse.ArgumentParser(description='MindSpore MNIST Example')
-parser.add_argument('--device_target', type=str, default="Ascend",
-                    choices=['Ascend', 'GPU'],
-                    help='device where the code will be implemented (default: Ascend)')
-parser.add_argument('--ckpt_path', type=str, default="",
-                    help='if mode is test, must provide path where the trained ckpt file')
+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="lenet", 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")
 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, device_target=args.device_target)
 
     # define fusion network
     network = LeNet5(cfg.num_classes)
     # load network checkpoint
-    param_dict = load_checkpoint(args.ckpt_path)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(network, param_dict)
 
     # export network
-    inputs = Tensor(np.ones([1, 1, cfg.image_height, cfg.image_width]), mindspore.float32)
-    export(network, inputs, file_name=cfg.air_name, file_format='AIR')
+    inputs = Tensor(np.ones([args.batch_size, 1, cfg.image_height, cfg.image_width]), mindspore.float32)
+    export(network, inputs, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/maskrcnn/export.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""export checkpoint file into air models"""
+"""export checkpoint file into air, onnx, mindir models"""
 import argparse
 import numpy as np
 
@@ -21,26 +21,32 @@ from mindspore import Tensor, context, load_checkpoint, load_param_into_net, exp
 from src.maskrcnn.mask_rcnn_r50 import Mask_Rcnn_Resnet50
 from src.config import config
 
-context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
+parser = argparse.ArgumentParser(description='maskrcnn 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="maskrcnn", 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()
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='maskrcnn_export')
-    parser.add_argument('--ckpt_file', type=str, default='', help='maskrcnn ckpt file.')
-    parser.add_argument('--output_file', type=str, default='', help='maskrcnn output air name.')
-    args_opt = parser.parse_args()
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
 
+if __name__ == '__main__':
     net = Mask_Rcnn_Resnet50(config=config)
-    param_dict = load_checkpoint(args_opt.ckpt_file)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(net, param_dict)
     net.set_train(False)
 
     bs = config.test_batch_size
 
-    img = Tensor(np.zeros([bs, 3, 768, 1280], np.float16))
-    img_metas = Tensor(np.zeros([bs, 4], np.float16))
-    gt_bboxes = Tensor(np.zeros([bs, 128, 4], np.float16))
-    gt_labels = Tensor(np.zeros([bs, 128], np.int32))
-    gt_num = Tensor(np.zeros([bs, 128], np.bool))
-    gt_mask = Tensor(np.zeros([bs, 128], np.bool))
-    export(net, img, img_metas, gt_bboxes, gt_labels, gt_num, gt_mask, file_name=args_opt.output_file,
-           file_format="AIR")
+    img = Tensor(np.zeros([args.batch_size, 3, config.img_height, config.img_width], np.float16))
+    img_metas = Tensor(np.zeros([args.batch_size, 4], np.float16))
+    gt_bboxes = Tensor(np.zeros([args.batch_size, config.num_gts, 4], np.float16))
+    gt_labels = Tensor(np.zeros([args.batch_size, config.num_gts], np.int32))
+    gt_num = Tensor(np.zeros([args.batch_size, config.num_gts], np.bool))
+    gt_mask = Tensor(np.zeros([args.batch_size, config.num_gts], np.bool))
+
+    input_data = [img, img_metas, gt_bboxes, gt_labels, gt_num, gt_mask]
+    export(net, *input_data, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/mobilenetv1/export.py

@@ -23,16 +23,17 @@ from src.mobilenet_v1 import mobilenet_v1 as mobilenet
 
 parser = argparse.ArgumentParser(description="mobilenetv1 export")
 parser.add_argument("--device_id", type=int, default=0, help="Device id")
-parser.add_argument("--batch_size", type=int, default=256, help="batch size")
 parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
 parser.add_argument("--dataset", type=str, default="imagenet2012", help="Dataset, either cifar10 or imagenet2012")
+parser.add_argument('--width', type=int, default=224, help='input width')
+parser.add_argument('--height', type=int, default=224, help='input height')
 parser.add_argument("--file_name", type=str, default="mobilenetv1", 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")
 args = parser.parse_args()
 
-context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args.device_id)
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
 
 if args.dataset == "cifar10":
     from src.config import config1 as config
@@ -40,8 +41,6 @@ else:
     from src.config import config2 as config
 
 if __name__ == "__main__":
-    config.batch_size = args.batch_size
-
     target = args.device_target
     if target != "GPU":
         context.set_context(device_id=args.device_id)
@@ -52,6 +51,6 @@ if __name__ == "__main__":
 
     network.set_train(False)
 
-    input_data = Tensor(np.zeros([config.batch_size, 3, 224, 224]).astype(np.float32))
+    input_data = Tensor(np.zeros([config.batch_size, 3, args.height, args.width]).astype(np.float32))
 
     export(network, input_data, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/mobilenetv2/export.py

@@ -13,22 +13,35 @@
 # limitations under the License.
 # ============================================================================
 """
-mobilenetv2 export mindir.
+mobilenetv2 export file.
 """
+import argparse
 import numpy as np
-from mindspore import Tensor, export
+from mindspore import Tensor, export, context
 from src.config import set_config
-from src.args import export_parse_args
 from src.models import define_net, load_ckpt
 from src.utils import set_context
 
+parser = argparse.ArgumentParser(description="mobilenetv2 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="mobilenetv2", help="output file name.")
+parser.add_argument("--file_format", type=str, choices=["AIR", "ONNX", "MINDIR"], default="AIR", help="file format")
+parser.add_argument('--platform', type=str, default="Ascend", choices=("Ascend", "GPU", "CPU"),
+                    help='run platform, only support GPU, CPU and Ascend')
+args = parser.parse_args()
+args.is_training = False
+args.run_distribute = False
+
+context.set_context(mode=context.GRAPH_MODE, device_target=args.platform, device_id=args.device_id)
+
 if __name__ == '__main__':
-    args_opt = export_parse_args()
-    cfg = set_config(args_opt)
+    cfg = set_config(args)
     set_context(cfg)
-    _, _, net = define_net(cfg, args_opt.is_training)
+    _, _, net = define_net(cfg, args.is_training)
 
-    load_ckpt(net, args_opt.pretrain_ckpt)
-    input_shp = [1, 3, cfg.image_height, cfg.image_width]
+    load_ckpt(net, args.ckpt_file)
+    input_shp = [args.batch_size, 3, cfg.image_height, cfg.image_width]
     input_array = Tensor(np.random.uniform(-1.0, 1.0, size=input_shp).astype(np.float32))
-    export(net, input_array, file_name=cfg.export_file, file_format=cfg.export_format)
+    export(net, input_array, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/nasnet/export.py

@@ -19,20 +19,27 @@ import argparse
 import numpy as np
 
 import mindspore as ms
-from mindspore import Tensor, load_checkpoint, load_param_into_net, export
+from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
 
 from src.config import nasnet_a_mobile_config_gpu as cfg
 from src.nasnet_a_mobile import NASNetAMobile
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='checkpoint export')
-    parser.add_argument('--checkpoint', type=str, default='', help='checkpoint of nasnet_a_mobile (Default: None)')
-    args_opt = parser.parse_args()
+parser = argparse.ArgumentParser(description='nasnet 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="nasnet", 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")
+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 = NASNetAMobile(num_classes=cfg.num_classes, is_training=False)
-    param_dict = load_checkpoint(args_opt.checkpoint)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(net, param_dict)
 
-    input_arr = Tensor(np.random.uniform(0.0, 1.0, size=[1, 3, cfg.image_size, cfg.image_size]), ms.float32)
-    export(net, input_arr, file_name=cfg.onnx_filename, file_format="ONNX")
-    export(net, input_arr, file_name=cfg.geir_filename, file_format="GEIR")
+    input_arr = Tensor(np.ones([args.batch_size, 3, cfg.image_size, cfg.image_size]), ms.float32)
+    export(net, input_arr, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/openpose/export.py

@@ -21,18 +21,27 @@ from mindspore import context
 from mindspore.train.serialization import load_checkpoint, load_param_into_net, export
 
 from src.openposenet import OpenPoseNet
+from src.config import params
 
-parser = argparse.ArgumentParser(description='checkpoint export')
-parser.add_argument('--checkpoint_path', type=str, default=None, help='Checkpoint file path')
-args_opt = parser.parse_args()
+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="openpose", 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()
 
-if __name__ == '__main__':
+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 = OpenPoseNet()
 
     # load checkpoint
-    param_dict = load_checkpoint(args_opt.checkpoint_path)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(net, param_dict)
-    inputs = np.random.uniform(0.0, 1.0, size=[1, 3, 368, 368]).astype(np.float32)
-    export(net, Tensor(inputs), file_name="openpose.air", file_format='AIR')
+    inputs = np.ones([args.batch_size, 3, params["insize"], params["insize"]]).astype(np.float32)
+    export(net, Tensor(inputs), file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/psenet/export.py

@@ -19,20 +19,27 @@ import argparse
 import numpy as np
 
 import mindspore as ms
-from mindspore import Tensor, load_checkpoint, load_param_into_net, export
+from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
 
 from src.config import config
 from src.ETSNET.etsnet import ETSNet
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='checkpoint export to air')
-    parser.add_argument('--checkpoint', type=str, default='', help='checkpoint of psenet (Default: None)')
-    args_opt = parser.parse_args()
+parser = argparse.ArgumentParser(description="psenet 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="psenet", 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__':
     net = ETSNet(config)
-    param_dict = load_checkpoint(args_opt.checkpoint)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(net, param_dict)
 
-    input_arr = Tensor(np.random.uniform(0.0, 1.0, size=[1, 3, config.INFER_LONG_SIZE, config.INFER_LONG_SIZE]),
-                       ms.float32)
-    export(net, input_arr, file_name=config.air_filename)
+    input_arr = Tensor(np.ones([args.batch_size, 3, config.INFER_LONG_SIZE, config.INFER_LONG_SIZE]), ms.float32)
+    export(net, input_arr, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/resnet/export.py

@@ -19,29 +19,40 @@ python export.py
 import argparse
 import numpy as np
 
-from mindspore import Tensor, load_checkpoint, load_param_into_net, export
+from mindspore import Tensor, load_checkpoint, load_param_into_net, export, context
+
+parser = argparse.ArgumentParser(description='resnet export')
+parser.add_argument('--network_dataset', type=str, default='resnet50_cifar10', choices=['resnet50_cifar10',
+                                                                                        'resnet50_imagenet2012',
+                                                                                        'resnet101_imagenet2012',
+                                                                                        "se-resnet50_imagenet2012"],
+                    help='network and dataset name.')
+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="resnet", help="output file name.")
+parser.add_argument('--width', type=int, default=224, help='input width')
+parser.add_argument('--height', type=int, default=224, help='input height')
+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__':
-    parser = argparse.ArgumentParser(description='resnet export')
-    parser.add_argument('--network_dataset', type=str, default='resnet50_cifar10', choices=['resnet50_cifar10',
-                                                                                            'resnet50_imagenet2012',
-                                                                                            'resnet101_imagenet2012',
-                                                                                            "se-resnet50_imagenet2012"],
-                        help='network and dataset name.')
-    parser.add_argument('--ckpt_file', type=str, default='', help='resnet ckpt file.')
-    parser.add_argument('--output_file', type=str, default='', help='resnet output air name.')
-    args_opt = parser.parse_args()
-
-    if args_opt.network_dataset == 'resnet50_cifar10':
+
+
+    if args.network_dataset == 'resnet50_cifar10':
         from src.config import config1 as config
         from src.resnet import resnet50 as resnet
-    elif args_opt.network_dataset == 'resnet50_imagenet2012':
+    elif args.network_dataset == 'resnet50_imagenet2012':
         from src.config import config2 as config
         from src.resnet import resnet50 as resnet
-    elif args_opt.network_dataset == 'resnet101_imagenet2012':
+    elif args.network_dataset == 'resnet101_imagenet2012':
         from src.config import config3 as config
         from src.resnet import resnet101 as resnet
-    elif args_opt.network_dataset == 'se-resnet50_imagenet2012':
+    elif args.network_dataset == 'se-resnet50_imagenet2012':
         from src.config import config4 as config
         from src.resnet import se_resnet50 as resnet
     else:
@@ -49,10 +60,10 @@ if __name__ == '__main__':
 
     net = resnet(config.class_num)
 
-    assert args_opt.ckpt_file is not None, "checkpoint_path is None."
+    assert args.ckpt_file is not None, "checkpoint_path is None."
 
-    param_dict = load_checkpoint(args_opt.ckpt_file)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(net, param_dict)
 
-    input_arr = Tensor(np.zeros([1, 3, 224, 224], np.float32))
-    export(net, input_arr, file_name=args_opt.output_file, file_format="AIR")
+    input_arr = Tensor(np.zeros([args.batch_size, 3, args.height, args.width], np.float32))
+    export(net, input_arr, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/resnet50_quant/export.py

@@ -1,51 +0,0 @@
-# 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 argparse
-import numpy as np
-
-from mindspore import context, Tensor
-from mindspore.train.serialization import export, load_checkpoint, load_param_into_net
-from mindspore.compression.quant import QuantizationAwareTraining
-
-from src.config import config_quant
-from modelsresnet_quant_manual import resnet50_quant
-
-parser = argparse.ArgumentParser(description='resnet50_quant 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("--img_size", type=int, default=224, help="image size")
-parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
-parser.add_argument("--file_name", type=str, default="resnet50_quant", help="output file name.")
-parser.add_argument('--file_format', type=str, choices=["AIR", "ONNX", "MINDIR"], default='MINDIR', 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__":
-    config = config_quant
-
-    network = resnet50_quant(class_num=config.class_num)
-
-    quantizer = QuantizationAwareTraining(bn_fold=True, per_channel=[True, False], symmetric=[True, False])
-    network = quantizer.quantize(network)
-
-    param_dict = load_checkpoint(args.ckpt_file)
-    load_param_into_net(network, param_dict)
-
-    network.set_train(False)
-    shape = [config.batch_size, 3] + [args.img_size, args.img_size]
-    input_data = Tensor(np.zeros(shape).astype(np.float32))
-
-    export(network, input_data, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/resnet_thor/export.py

@@ -21,9 +21,18 @@ from src.resnet_thor import resnet50 as resnet
 from src.config import config
 
 parser = argparse.ArgumentParser(description='checkpoint export')
-parser.add_argument('--checkpoint_path', type=str, default=None, help='Checkpoint file path')
-parser.add_argument('--output_file', type=str, default='', help='resnet output air name.')
-args_opt = parser.parse_args()
+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('--width', type=int, default=224, help='input width')
+parser.add_argument('--height', type=int, default=224, help='input height')
+parser.add_argument("--file_name", type=str, default="resnet_thor", 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__':
 
@@ -34,12 +43,12 @@ if __name__ == '__main__':
     net.add_flags_recursive(thor=False)
 
     # load checkpoint
-    param_dict = load_checkpoint(args_opt.checkpoint_path)
+    param_dict = load_checkpoint(args.ckpt_file)
     keys = list(param_dict.keys())
     for key in keys:
         if "damping" in key:
             param_dict.pop(key)
     load_param_into_net(net, param_dict)
 
-    inputs = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32)
-    export(net, Tensor(inputs), file_name=args_opt.output_file, file_format='AIR')
+    inputs = np.random.uniform(0.0, 1.0, size=[args.batch_size, 3, args.height, args.width]).astype(np.float32)
+    export(net, Tensor(inputs), file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/resnext50/export.py

@@ -21,34 +21,25 @@ from mindspore import context, Tensor, load_checkpoint, load_param_into_net, exp
 from src.config import config
 from src.image_classification import get_network
 
-
-def parse_args():
-    """parse_args"""
-    parser = argparse.ArgumentParser('mindspore classification test')
-    parser.add_argument('--platform', type=str, default='Ascend', choices=('Ascend', 'GPU'), help='run platform')
-
-    parser.add_argument('--pretrained', type=str, required=True, help='fully path of pretrained model to load. '
-                        'If it is a direction, it will test all ckpt')
-
-    args, _ = parser.parse_known_args()
-    args.image_size = config.image_size
-    args.num_classes = config.num_classes
-
-    args.image_size = list(map(int, config.image_size.split(',')))
-    args.image_height = args.image_size[0]
-    args.image_width = args.image_size[1]
-    args.export_format = config.export_format
-    args.export_file = config.export_file
-    return args
+parser = argparse.ArgumentParser(description='checkpoint 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('--width', type=int, default=224, help='input width')
+parser.add_argument('--height', type=int, default=224, help='input height')
+parser.add_argument("--file_name", type=str, default="resnext50", 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__':
-    args_export = parse_args()
-    context.set_context(mode=context.GRAPH_MODE, device_target=args_export.platform)
-
-    net = get_network(num_classes=args_export.num_classes, platform=args_export.platform)
+    net = get_network(num_classes=config.num_classes, platform=args.device_target)
 
-    param_dict = load_checkpoint(args_export.pretrained)
+    param_dict = load_checkpoint(args.ckpt_file)
     load_param_into_net(net, param_dict)
-    input_shp = [1, 3, args_export.image_height, args_export.image_width]
+    input_shp = [args.batch_size, 3, args.height, args.width]
     input_array = Tensor(np.random.uniform(-1.0, 1.0, size=input_shp).astype(np.float32))
-    export(net, input_array, file_name=args_export.export_file, file_format=args_export.export_format)
+    export(net, input_array, file_name=args.file_name, file_format=args.file_format)

model_zoo/official/cv/shufflenetv2/export.py

@@ -0,0 +1,49 @@
+# 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.
+# ============================================================================
+"""evaluate_imagenet"""
+import argparse
+import numpy as np
+
+import mindspore as ms
+from mindspore import context, Tensor, load_checkpoint, load_param_into_net, export
+
+from src.config import config_gpu as cfg
+from src.shufflenetv2 import ShuffleNetV2
+
+parser = argparse.ArgumentParser(description='checkpoint export')
+parser.add_argument("--device_id", type=int, default=0, help="Device id")
+parser.add_argument("--batch_size", type=int, default=128, help="batch size")
+parser.add_argument("--ckpt_file", type=str, required=True, help="Checkpoint file path.")
+parser.add_argument('--width', type=int, default=224, help='input width')
+parser.add_argument('--height', type=int, default=224, help='input height')
+parser.add_argument("--file_name", type=str, default="shufflenetv2", 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="GPU",
+                    choices=["Ascend", "GPU", "CPU"], help="device where the code will be implemented (default: GPU)")
+args = parser.parse_args()
+
+context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id)
+
+if __name__ == '__main__':
+    if args.device_target != 'GPU':
+        raise ValueError("Only supported GPU now.")
+
+    net = ShuffleNetV2(n_class=cfg.num_classes)
+    ckpt = load_checkpoint(args.ckpt_file)
+    load_param_into_net(net, ckpt)
+    net.set_train(False)
+
+    input_data = Tensor(np.ones([args.batch_size, 3, args.height, args.width]), ms.float32)
+    export(net, input_data, file_name=args.file_name, file_format=args.file_format)