p60975123
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!11541 modify transformer export 

From: @changzherui
Reviewed-by: @guoqi1024,@kingxian
Signed-off-by: @guoqi1024
This commit is contained in:
mindspore-ci-bot 2021-01-25 10:28:57 +08:00 committed by Gitee
parent 54b8d53780 a0ff6613c3
commit 14a6713d08
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
model_zoo/official/nlp/transformer/export.py
View File

@ -26,7 +26,6 @@ from eval import load_weights
parser = argparse.ArgumentParser(description='transformer 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("--file_name", type=str, default="transformer", 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",
@ -43,7 +42,7 @@ if __name__ == '__main__':
parameter_dict = load_weights(cfg.model_file)
load_param_into_net(tfm_model, parameter_dict)
source_ids = Tensor(np.ones((args.batch_size, transformer_net_cfg.seq_length)).astype(np.int32))
source_mask = Tensor(np.ones((args.batch_size, transformer_net_cfg.seq_length)).astype(np.int32))
source_ids = Tensor(np.ones((transformer_net_cfg.batch_size, transformer_net_cfg.seq_length)).astype(np.int32))
source_mask = Tensor(np.ones((transformer_net_cfg.batch_size, transformer_net_cfg.seq_length)).astype(np.int32))
export(tfm_model, source_ids, source_mask, file_name=args.file_name, file_format=args.file_format)