def save_model(
args,
model,
optimizer,
src_tokenizer: BertTokenizer,
tgt_tokenizer: GPT2Tokenizer,
nstep,
nepoch,
bleu,
loss,
):
# 记录整体训练评价结果
train_metric_log_file = os.path.join(args.output_dir, "training_metric.tsv")
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
if os.path.exists(train_metric_log_file):
with open(train_metric_log_file, "a", encoding="utf-8") as fa:
fa.write("{}\t{}\t{}\t{}\n".format(nepoch, nstep, loss, bleu))
else:
with open(train_metric_log_file, "w", encoding="utf-8") as fw:
fw.write("epoch\tstep\tloss\tbleu\n")
fw.write("{}\t{}\t{}\t{}\n".format(nepoch, nstep, loss, bleu))
# 保存模型
model_save_path = os.path.join(
args.output_dir, "epoch{}_step{}/".format(nepoch, nstep)
)
os.makedirs(model_save_path)
model.save_pretrained(model_save_path)
if local_rank == 0 or local_rank == -1:
print(
"epoch:{} step:{} loss:{} bleu:{} model save complete.".format(
nepoch, nstep, round(loss, 4), round(bleu, 4)
)
)
if args.save_optimizer:
torch.save(optimizer, os.path.join(model_save_path, "optimizer.pt"))
# 保存tokenizer
src_tokenizer.save_pretrained(os.path.join(model_save_path, "src_tokenizer"))
tgt_tokenizer.save_pretrained(os.path.join(model_save_path, "tgt_tokenizer"))
