PyTorch Lightning从Checkpoint恢复训练的状态恢复疑问

问题背景 #

我正在运行一个已编写完成的PyTorch模型，原代码中作者使用PyTorch Lightning的Trainer保存每个epoch的ckpt文件，但未提供从checkpoint恢复训练的选项。

原代码 #

trainer = pl.Trainer.from_argparse_args(args,
                                        default_root_dir=args.logdir,
                                        gpus = args.gpus,
                                        accelerator='ddp',
                                        sync_batchnorm=True,
                                        plugins=DDPPlugin(find_unused_parameters=False),
                                        profiler='simple',
                                        benchmark=True,
                                        log_every_n_steps=1,
                                        flush_logs_every_n_steps=5,
                                        callbacks=[checkpoint_callback,
                                                    ],
                                        check_val_every_n_epoch = args.val_every,
                                        max_epochs = args.epochs,
                                        logger=logger
                                        )

修改后的代码（实现从指定checkpoint恢复训练） #

if args.loadfromcheckpoint > 0:
        trainer = pl.Trainer(
            resume_from_checkpoint=args.logdir+"/epoch={checkpoint}-last.ckpt".format(checkpoint=args.loadfromcheckpoint),
            default_root_dir=args.logdir,
            gpus = args.gpus,
            accelerator='ddp',
            sync_batchnorm=True,
            plugins=DDPPlugin(find_unused_parameters=False),
            profiler='simple',
            benchmark=True,
            log_every_n_steps=1,
            flush_logs_every_n_steps=5,
            callbacks=[checkpoint_callback,
                        ],
            check_val_every_n_epoch = args.val_every,
            max_epochs = args.epochs,
            logger=logger)
        trainer.fit(TCP_model, dataloader_train, dataloader_val)
    else:
        trainer.fit(TCP_model, dataloader_train, dataloader_val)

修改后的代码可正常运行，但存在以下两个技术疑问：

1. Lightning API仅恢复model的state_dict，还是会同时恢复optimizer_states、lr_schedulers等状态？
2. 若Lightning未加载这些状态，该如何手动加载？

解答 #

1. Lightning恢复checkpoint时包含的内容 #

当你使用resume_from_checkpoint参数初始化Trainer并调用fit时，Lightning会自动恢复模型的state_dict、优化器的状态、学习率调度器的状态，甚至包括训练的epoch数、全局步数等训练进度信息。

Lightning默认保存的ckpt文件是完整的训练快照，包含所有必要的状态数据，能确保你从断点处无缝继续训练，无需额外手动处理这些状态。

2. 手动加载状态的方法（若需要） #

如果因为自定义特殊组件、ckpt文件不完整等原因需要手动加载状态，可按以下步骤操作：

• 加载模型状态：

  TCP_model.load_state_dict(torch.load(checkpoint_path)['state_dict'])
  

• 加载优化器状态：
  先按训练时的配置初始化优化器，再加载状态：

  optimizer = torch.optim.Adam(TCP_model.parameters(), lr=1e-3)
  optimizer.load_state_dict(torch.load(checkpoint_path)['optimizer_states'][0])
  

  注：optimizer_states是列表格式，单个优化器时取索引0即可，多优化器场景按对应索引取值。

• 加载学习率调度器状态：
  先初始化调度器，再加载状态：

  scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10)
  scheduler.load_state_dict(torch.load(checkpoint_path)['lr_schedulers'][0])
  

完成上述加载后，无需在Trainer初始化时指定resume_from_checkpoint，只需将模型、优化器、调度器传入LightningModule的configure_optimizers方法，确保训练时使用这些已加载状态的组件即可。

内容的提问来源于stack exchange，提问作者Dinura Dissanayake