!9704 fix export file of resnet_thor

From: @wangmin0104
Reviewed-by: @sunnybeike,@wang_zi_dong
Signed-off-by: @sunnybeike
This commit is contained in:
mindspore-ci-bot 2020-12-10 09:39:08 +08:00 committed by Gitee
commit d1af8d3693
1 changed files with 2 additions and 1 deletions

View File

@ -22,6 +22,7 @@ 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()
if __name__ == '__main__':
@ -41,4 +42,4 @@ if __name__ == '__main__':
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='resnet-42_5004', file_format='AIR')
export(net, Tensor(inputs), file_name=args_opt.output_file, file_format='AIR')