PyTorch Lightning如何从最后epoch及权重恢复训练状态

PyTorch Lightning 原生提供完整的训练状态自动恢复能力，不需要手动重载回调、解析CSV日志实现恢复逻辑，你之前加载权重时丢失学习率、epoch计数等信息，本质是使用了纯权重加载方式，没有走框架原生的断点续训逻辑。

• 原生checkpoint的存储范围
  框架自带的ModelCheckpoint回调生成的.ckpt文件，不是仅存储模型权重，会全量序列化所有训练上下文：包括当前epoch编号、全局训练step数、优化器全量状态（含当前学习率、动量等超参数实时值）、学习率调度器状态、所有回调的内部状态、甚至训练过程的随机数状态，完全满足断点续训的状态一致性要求。
• 自动恢复的使用方式
  不需要手动调用load_state_dict加载权重，初始化Trainer时直接传入ckpt_path参数指向目标ckpt文件即可，框架会在fit启动时自动恢复所有状态，训练计数、学习率调度逻辑会和中断前完全对齐。示例代码：

  from pytorch_lightning import Trainer
  from pytorch_lightning.callbacks import ModelCheckpoint
  
  # 开启last checkpoint保存，每次训练epoch结束自动更新最新断点
  ckpt_callback = ModelCheckpoint(dirpath="./train_ckpts/", save_last=True)
  # 传入ckpt_path即可触发全量状态自动恢复
  trainer = Trainer(
      callbacks=[ckpt_callback],
      ckpt_path="./train_ckpts/last.ckpt",
      max_epochs=100
  )
  trainer.fit(your_model, your_dataloader)
  

注意：如果你之前是手动调用torch.save(model.state_dict(), "xxx.pt")存储的纯权重文件，文件本身不包含训练状态信息，这种场景才需要手动补充恢复逻辑。

• 常见的状态丢失原因
  如果你用.ckpt文件加载还是出现状态丢失，基本是因为没有通过Trainer的ckpt_path参数传入断点，而是手动在代码里调用model.load_state_dict(torch.load(ckpt_path)["state_dict"])加载权重，这种方式只会初始化模型参数，不会加载优化器、调度器、训练计数等上下文，本质是用旧权重做初始化，不属于断点续训。
• 不需要自定义回调或解析日志
  只要使用框架原生的ModelCheckpoint存储断点，ckpt_path参数会自动处理所有状态恢复逻辑，不需要重载Checkpoint组件，也不需要读取CSVLogger的日志文件回溯最后epoch编号。

内容的提问来源于stack exchange，提问作者Mohbat Tharani