PyTorch Lightning中last.ckpt恢复训练问题的分析与解决

问题背景

在使用PyTorch Lightning进行模型训练时，开发者经常会利用ModelCheckpoint回调来保存训练过程中的检查点。其中，save_last参数被设计用来保存一个指向最新检查点的符号链接(last.ckpt)，方便开发者随时恢复训练。

然而，在PyTorch Lightning 2.1.x版本中，存在一个bug导致无法从last.ckpt恢复训练。当尝试使用trainer.fit(ckpt_path="path/to/last.ckpt")时，系统会抛出"Checkpoint file not found"错误，尽管文件确实存在且符号链接也正确指向了有效的检查点文件。

问题复现

通过以下代码可以复现该问题：

import lightning.pytorch as pl
from lightning.pytorch import LightningModule, Trainer
from lightning.pytorch.callbacks import ModelCheckpoint
from torch.utils.data import DataLoader, Dataset

# 定义简单数据集和模型
class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.data = torch.randn(length, size)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
    
    def forward(self, x):
        return self.layer(x)
    
    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return loss
    
    def configure_optimizers(self):
        return torch.optim.SGD(self.parameters(), lr=0.1)

# 配置训练参数
trainer_args = {
    "max_epochs": 10,
    "callbacks": [ModelCheckpoint(monitor="train_loss", save_last="link", mode="min", save_top_k=5)]
}

# 第一次训练
train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
model = SimpleModel()
trainer = Trainer(**trainer_args)
trainer.fit(model, train_data)

# 尝试从last.ckpt恢复训练
trainer = Trainer(**trainer_args)
trainer.fit(model, train_data, ckpt_path="lightning_logs/version_0/checkpoints/last.ckpt")

问题分析

该问题主要源于PyTorch Lightning 2.1.x版本中检查点加载逻辑的一个缺陷。当使用符号链接作为检查点路径时，系统无法正确解析符号链接指向的实际文件路径，导致文件查找失败。

具体表现为：

1. save_last="link"参数正确创建了last.ckpt符号链接
2. 符号链接正确指向了最新的检查点文件
3. 但在尝试加载时，系统无法处理符号链接，直接尝试加载符号链接本身而非其指向的文件

解决方案

该问题已在PyTorch Lightning的代码库中得到修复，修复提交为#19303。修复后的版本将正确处理符号链接，能够成功从last.ckpt恢复训练。

对于遇到此问题的开发者，建议采取以下措施：

1. 升级到PyTorch Lightning 2.1.4或更高版本
2. 如果暂时无法升级，可以手动解析符号链接路径：

import os

# 手动解析符号链接
real_ckpt_path = os.path.realpath("lightning_logs/version_0/checkpoints/last.ckpt")
trainer.fit(model, train_data, ckpt_path=real_ckpt_path)

最佳实践

在使用PyTorch Lightning的检查点功能时，建议：

1. 明确指定save_last参数为True或"link"
2. 定期检查保存的检查点文件是否完整
3. 在恢复训练前，先验证检查点文件是否可加载
4. 保持PyTorch Lightning版本更新，以获取最新的bug修复和功能改进

通过理解并正确处理检查点恢复问题，开发者可以更高效地利用PyTorch Lightning进行模型训练，特别是在需要中断后恢复训练的场景下。