pytorch 模型保存和断点恢复训练

           在使用pytorch训练模型中，有时候需要从断点处继续训练，那么需要将模型、优化器、lr_scheduler,epoch和其他args进行保存，如果使用的apm混合精度，还要保存scaler。在恢复训练的时候，需要重新加载数据即可,在推理阶段，只要保存model.state_dict()就好了。

checkpoint  = {
                "model" : model.state_dict(),
                "optimizer":optimizer.state_dict(),
                "lr_scheduler": lr_scheduler.state_dict(),
                "args":args  # 其他参数
                "epoch": epoch,

}

if args.amp:
checkpoint["scaler] = scaler.state_dict()

torch.save(checkpoint,"file_path")

# 加载模型

checkpoint = torch.load("file_path",map_location="cpu")
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
epoch = checkpoint["lr_scheduler"]+1
if amp:
    scaler.load_state_dict(checkpoint["scaler"])