模型权重文件的保存与加载

一般地，在训练过程中有可能会因为某些原因导致训练过程的终止，尤其是在训练一些大型模型的时候，所以在训练过程中保存权重文件是很有必要的，要不然训练好几天的代码又要重新跑了

下面来讲一下通常权重文件中都保存一些什么内容？

checkpoint = {
    'epoch'：epoch                           # 当前迭代的轮次
    "model":model.state_dict(),  # 模型权重
    "optimizer":optimizer.state_dict(),      # 优化器的参数
    "lr":lr                                  # 学习率策略
    "lr_scheduler":lr_scheduler.state_dict(),
    "args":args

}
if args.amp:   # 混合精度训练
    checkpoint["scaler"] = scaler.state_dict()

权重保存方式	存储大小
仅保留模型权重	169MB
包含优化器等信息	339MB

仅保留模型权重：

save_files={
    'model':model.state_dict(),
}

包含优化器等信息：

save_files = {
    'model':model.state_dict(),
    'optimizer':optimizer.state_dict(),
    'lr_scheduler':lr_scheduler.state_dict(),
    'epoch':epoch
}

所以保存的权重文件一般都很大，但是我们在载入预训练权重的时候只需要载入模型的参数，即只载入 model.state_dict()

如何正确载入上一次的模型权重继续进行训练？

if args.resume:
    checkpoint = torch.load(args.resume,map_location="cpu")
    model.load_state_dict(checkpoint['model'])   # 载入模型的参数
    optimizer.load_state_dict(checkpoint["optimizer"])  # 载入优化器的相关信息
    lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
    args.start_epoch = checkpoint["epoch"]+1
    if args.amp:
        scaler.load_state_dict(checkpoint["scaler"])

 如果只加载模型权重相关信息的话，继续训练的结果其实并不能对齐

比如在优化器当中，如果使用了

momentum

动量，会保存在训练过程中的统计信息，继续训练，没有载入优化器相关信息的话就会丢失之前的统计信息，这样就不算真正的继续训练

比如对于学习率而言，在训练多个epoch后，学习率 lr 可能已经衰减到某一个值了，如果没有载入历史的学习率信息的话，学习率又是从初始的学习率开始降低，而不是上次训练结束的学习率

综上所述，

• 为了训练过程中保留更多有用信息，除了保存模型权重外还会保存优化器等信息
• 如果用权重文件进行预测操作的时候，不再需要其它信息，可载入模型后重新保存权重 （仅模型权重）