dev_info = processor.get_examples(dev_raw_examples, 'dev')
tmp_output_dir = os.path.join(base_output_dir, f'v{i}')
opt.output_dir = tmp_output_dir
train_base(opt, train_examples, dev_info)
if __name__ == '__main__':
args = Args().get_parser()
assert args.task_type in ['crf', 'span', 'mrc']
args.output_dir = os.path.join(args.output_dir, args.task_type)
set_seed(args.seed)
if args.attack_train != '':
args.output_dir += f'_{args.attack_train}'
if args.weight_decay:
args.output_dir += '_wd'
if args.use_fp16:
args.output_dir += '_fp16'
if args.task_type == 'span':
args.output_dir += f'_{args.loss_type}'
if args.task_type == 'mrc':
args.output_dir += f'_{args.loss_type}'
scaler.step(optimizer)
scaler.update()
else:
optimizer.step()
if batch % 100:
torch.cuda.empty_cache()
# if config.local_rank in [0, -1]:
now = strftime("%Y-%m-%d %H:%M:%S", localtime())
print("time:{},epoch:{}/{},loss:{}".format(now, epoch + 1,
config.num_train_epochs,
loss.item()))
if torch.cuda.device_count() > 1:
checkpoint = model.module.state_dict()
else:
checkpoint = model.state_dict()
torch.save(checkpoint,
'./data/pre_bert/nezha-base-www/pytorch_model.bin')
del checkpoint
if __name__ == '__main__':
print("模型预训练开始：{}".format(
time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime())))
config = bert_config()
set_seed(config.seed)
pre_trained(config)
print("模型预训练结束：{}".format(
time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime())))