diff --git a/model_zoo/official/cv/alexnet/export.py b/model_zoo/official/cv/alexnet/export.py index f7d08786ad2..f145abe91f3 100644 --- a/model_zoo/official/cv/alexnet/export.py +++ b/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) diff --git a/model_zoo/official/cv/centerface/export.py b/model_zoo/official/cv/centerface/export.py index 0f46f1b83d9..a54fdb83703 100644 --- a/model_zoo/official/cv/centerface/export.py +++ b/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() diff --git a/model_zoo/official/cv/cnnctc/export.py b/model_zoo/official/cv/cnnctc/export.py index a8f7a4f781b..30ed12a756c 100644 --- a/model_zoo/official/cv/cnnctc/export.py +++ b/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) diff --git a/model_zoo/official/cv/deeplabv3/export.py b/model_zoo/official/cv/deeplabv3/export.py index da17dc3936b..fe436023fa5 100644 --- a/model_zoo/official/cv/deeplabv3/export.py +++ b/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() + +context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id) 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() - 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) diff --git a/model_zoo/official/cv/densenet121/export.py b/model_zoo/official/cv/densenet121/export.py index 329f5b084a3..ee95b7b18ae 100644 --- a/model_zoo/official/cv/densenet121/export.py +++ b/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) diff --git a/model_zoo/official/cv/faster_rcnn/export.py b/model_zoo/official/cv/faster_rcnn/export.py index 999868bb788..36ed88bf562 100644 --- a/model_zoo/official/cv/faster_rcnn/export.py +++ b/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) diff --git a/model_zoo/official/cv/googlenet/export.py b/model_zoo/official/cv/googlenet/export.py index c911ac2ce20..28b260bd056 100644 --- a/model_zoo/official/cv/googlenet/export.py +++ b/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() - if args_opt.dataset_name == 'cifar10': +context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id) + +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) diff --git a/model_zoo/official/cv/inceptionv3/export.py b/model_zoo/official/cv/inceptionv3/export.py index 61daaaa9751..9fb06a35a9d 100644 --- a/model_zoo/official/cv/inceptionv3/export.py +++ b/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) diff --git a/model_zoo/official/cv/inceptionv4/export.py b/model_zoo/official/cv/inceptionv4/export.py index 6b0fe4d4b13..8b51c824605 100644 --- a/model_zoo/official/cv/inceptionv4/export.py +++ b/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() + +context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id) if __name__ == '__main__': - - args_opt = parse_args() - 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) diff --git a/model_zoo/official/cv/lenet/export.py b/model_zoo/official/cv/lenet/export.py index 1f8ddf56e34..efff69db6cf 100644 --- a/model_zoo/official/cv/lenet/export.py +++ b/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) diff --git a/model_zoo/official/cv/maskrcnn/export.py b/model_zoo/official/cv/maskrcnn/export.py index 14ee325be91..e321c29e601 100644 --- a/model_zoo/official/cv/maskrcnn/export.py +++ b/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() + +context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, device_id=args.device_id) 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() - 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) diff --git a/model_zoo/official/cv/mobilenetv1/export.py b/model_zoo/official/cv/mobilenetv1/export.py index 9af3d3cf0a7..4b5264deec5 100644 --- a/model_zoo/official/cv/mobilenetv1/export.py +++ b/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) diff --git a/model_zoo/official/cv/mobilenetv2/export.py b/model_zoo/official/cv/mobilenetv2/export.py index e9cf7f9d3c3..037a1294d0b 100644 --- a/model_zoo/official/cv/mobilenetv2/export.py +++ b/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 -if __name__ == '__main__': - args_opt = export_parse_args() - cfg = set_config(args_opt) - set_context(cfg) - _, _, net = define_net(cfg, args_opt.is_training) +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 - load_ckpt(net, args_opt.pretrain_ckpt) - input_shp = [1, 3, cfg.image_height, cfg.image_width] +context.set_context(mode=context.GRAPH_MODE, device_target=args.platform, device_id=args.device_id) + +if __name__ == '__main__': + cfg = set_config(args) + set_context(cfg) + _, _, net = define_net(cfg, args.is_training) + + 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) diff --git a/model_zoo/official/cv/nasnet/export.py b/model_zoo/official/cv/nasnet/export.py index 73907bb2ffb..9dcf00e70a2 100755 --- a/model_zoo/official/cv/nasnet/export.py +++ b/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) diff --git a/model_zoo/official/cv/openpose/export.py b/model_zoo/official/cv/openpose/export.py index 7bf87275e7f..dbc00e8ca81 100644 --- a/model_zoo/official/cv/openpose/export.py +++ b/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) diff --git a/model_zoo/official/cv/psenet/export.py b/model_zoo/official/cv/psenet/export.py index 84ac1ab7c0a..3adfa109210 100755 --- a/model_zoo/official/cv/psenet/export.py +++ b/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) diff --git a/model_zoo/official/cv/resnet/export.py b/model_zoo/official/cv/resnet/export.py index d0b3866380b..2c3fe625543 100644 --- a/model_zoo/official/cv/resnet/export.py +++ b/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) diff --git a/model_zoo/official/cv/resnet50_quant/export.py b/model_zoo/official/cv/resnet50_quant/export.py deleted file mode 100644 index ff4f161b793..00000000000 --- a/model_zoo/official/cv/resnet50_quant/export.py +++ /dev/null @@ -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) diff --git a/model_zoo/official/cv/resnet_thor/export.py b/model_zoo/official/cv/resnet_thor/export.py index f90cfa0467f..06111b80645 100644 --- a/model_zoo/official/cv/resnet_thor/export.py +++ b/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) diff --git a/model_zoo/official/cv/resnext50/export.py b/model_zoo/official/cv/resnext50/export.py index b2dc045698e..e59aa5a876b 100644 --- a/model_zoo/official/cv/resnext50/export.py +++ b/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 +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() -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 +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=config.num_classes, platform=args.device_target) - net = get_network(num_classes=args_export.num_classes, platform=args_export.platform) - - 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) diff --git a/model_zoo/official/cv/shufflenetv2/export.py b/model_zoo/official/cv/shufflenetv2/export.py new file mode 100644 index 00000000000..840ed610120 --- /dev/null +++ b/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)