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

!5368 fix sink_size bug for transformer 

Merge pull request !5368 from yuchaojie/transformer_st
This commit is contained in:
mindspore-ci-bot 2020-08-28 11:12:43 +08:00 committed by Gitee
parent cd9505f5b6 9fb6f0c34b
commit f3fc24c82a
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
model_zoo/official/nlp/transformer/train.py
View File

@ -170,8 +170,14 @@ def run_transformer_train():
netwithgrads.set_train(True)
model = Model(netwithgrads)
model.train(args.epoch_size, dataset, callbacks=callbacks, dataset_sink_mode=(args.enable_data_sink == "true"),
sink_size=args.save_checkpoint_steps)
enable_sink = (args.enable_data_sink == "true")
if enable_sink:
sink_size = args.save_checkpoint_steps
model.train(args.epoch_size*dataset.get_dataset_size()//sink_size, dataset, callbacks=callbacks,
dataset_sink_mode=enable_sink, sink_size=sink_size)
else:
model.train(args.epoch_size, dataset, callbacks=callbacks, dataset_sink_mode=enable_sink)
if __name__ == '__main__':
run_transformer_train()