PyTorch Lightning每个Epoch后执行推理生成图表的最优实现方案

我正在使用PyTorch Lightning，想要在每个训练Epoch结束后，在小型数据集上执行推理并生成图表，通过Weights & Biases进行监控。原本打算用带有on_train_epoch_end方法的Callback来生成图表，在该方法中调用trainer.predict执行推理，但触发了如下AssertionError错误，求自然优雅的正确实现方式。

最小可复现示例 #

import lightning as L
from lightning.pytorch.callbacks import Callback

import torch
from torch.utils.data import DataLoader
from torch import nn, optim

class Model(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.f = nn.Linear(10, 1)
        
    def training_step(self, batch, *args):
        out = self(batch)
        return out.mean() ** 2
    
    def forward(self, x):
        return self.f(x)[:, 0]

    def train_dataloader(self):
        return DataLoader(torch.randn((100, 10)))
    
    def predict_dataloader(self):
        return DataLoader(torch.randn((100, 10)))
    
    def predict_step(self, batch):
        return self(batch)
    
    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=1e-3)
        return optimizer
    
class CallbackExample(Callback):
    def on_train_epoch_end(self, trainer: L.Trainer, model: Model) -> None:
        loader = model.predict_dataloader()
        trainer.predict(model, loader)
        
        ... # 将图表保存到wandb

model = Model()
callback = CallbackExample()
trainer = L.Trainer(max_epochs=2, callbacks=callback, accelerator="mps")

trainer.fit(model)

报错信息 #

File ~/Library/Caches/pypoetry/virtualenvs/novae-ezkWKrh6-py3.9/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py:233, in _LoggerConnector.metrics(self)
    231 """This function returns either batch or epoch metrics."""
    232 on_step = self._first_loop_iter is not None
--> 233 assert self.trainer._results is not None
    234 return self.trainer._results.metrics(on_step)

AssertionError:

问题根源 #

直接在on_train_epoch_end里调用trainer.predict会导致训练器内部状态冲突——此时训练循环尚未完全结束，训练器的结果对象（_results）还未初始化，从而触发断言错误。

正确实现方式 #

不要调用trainer.predict，而是手动遍历预测数据集，利用模型执行推理，同时确保模型处于评估模式，避免批量归一化、dropout等训练层影响结果：

修改后的Callback代码如下：

class CallbackExample(Callback):
    def on_train_epoch_end(self, trainer: L.Trainer, model: Model) -> None:
        # 切换模型到评估模式
        model.eval()
        # 获取预测数据集
        loader = model.predict_dataloader()
        predictions = []
        
        # 禁用梯度计算，节省内存和计算资源
        with torch.no_grad():
            for batch in loader:
                # 转移数据到模型所在设备
                batch = batch.to(model.device)
                # 执行推理
                pred = model.predict_step(batch)
                predictions.append(pred)
        
        # 合并所有预测结果
        predictions = torch.cat(predictions)
        
        # ---------------------------
        # 这里添加生成图表的代码，比如用matplotlib或seaborn
        # 然后将图表上传到Weights & Biases
        # 示例：
        # import wandb
        # import matplotlib.pyplot as plt
        # plt.hist(predictions.cpu().numpy(), bins=20)
        # wandb.log({"prediction_distribution": plt}, step=trainer.current_epoch)
        # plt.close()
        # ---------------------------
        
        # 切换回训练模式
        model.train()

额外优化建议 #

• 如果预测数据集很小，可以直接在CPU上处理，避免设备转移的开销
• 可以将预测数据集的定义单独提取出来，避免每次都重新创建DataLoader
• 若需要复用predict_step中的逻辑，手动调用是最安全的方式，不会干扰训练器的内部状态

内容的提问来源于stack exchange，提问作者Quentin BLAMPEY