调用trainer.tune后PyTorch Lightning的ModelCheckpoint找不到val_loss报错

在Kaggle环境下，使用PyTorch Lightning结合wandb logger训练LSTM模型时，遇到如下问题：

• 不执行trainer.tune()时，trainer.fit()可以正常运行
• 执行trainer.tune()后，ModelCheckpoint抛出错误，提示找不到val_loss指标，尽管已经在validation_step中明确log了该指标，且通过打印确认val_loss在ModelCheckpoint触发前已计算，但训练指标中仍缺失val_loss。

模型类代码 #

class Model(pl.LightningModule):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        bidirectional: bool = False,
        lstm_layers: int = 1,
        lstm_dropout: float = 0.4,
        fc_dropout: float = 0.4,
        lr: float = 0.01,
        lr_scheduler_patience: int = 2,
    ):
        super().__init__()
        self.lr = lr
        self.save_hyperparameters()

        # LSTM
        self.encoder_lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=lstm_layers,
            bidirectional=bidirectional,
            dropout=lstm_dropout if lstm_layers > 1 else 0,
            batch_first=True,
        )

        # Fully-connected
        num_directions = 2 if bidirectional else 1
        self.fc = nn.Sequential(
            nn.Linear(
                hidden_size * num_directions, hidden_size * num_directions * 2
            ),
            nn.ReLU(),
            nn.Dropout(fc_dropout),
            nn.Linear(hidden_size * num_directions * 2, input_size),
        )

        self.loss_function = nn.MSELoss()

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
                    optimizer, patience=self.hparams.lr_scheduler_patience
                ),
                "monitor": "val_loss",
            },
        }

    def forward(self, x, prev_state):
        ...

    def training_step(self, batch, batch_idx):
        loss, _ = self._step(batch)

        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        loss, embeddings = self._step(batch)

        self.log("val_loss", loss)
        
        return {
            'val_loss': loss,
            'preds': embeddings # 供自定义回调使用
        }

    def test_step(self, batch, batch_idx):
        loss, _ = self._step(batch)

        self.log("test_loss", loss)

训练代码 #

model = Model(
    bidirectional=False,
    lstm_layers=1,
    lstm_dropout=0.4,
    fc_dropout=0.4,
    lr=0.01,
    lr_scheduler_patience=2
)

...

checkpoint_callback = ModelCheckpoint(
    monitor="val_loss",
    every_n_train_steps=100,
    verbose=True
)

trainer = pl.Trainer(
    accelerator='gpu',
    precision=16,
    max_epochs=100,
    callbacks=[early_stopping, checkpoint_callback, lr_monitor, custom_callback],
    log_every_n_steps=50,
    logger=wandb_logger,
    auto_lr_find=True,
)

trainer.tune(model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)

trainer.fit(model, train_dataloader, val_dataloader)

报错信息 #

MisconfigurationException: `ModelCheckpoint(monitor='val_loss')` could not find the monitored key in the returned metrics: ['train_loss', 'epoch', 'step']. HINT: Did you call `log('val_loss', value)` in the `LightningModule`?

解决方案 #

1. 临时禁用tune阶段的ModelCheckpoint #

trainer.tune()会触发部分训练/验证流程，但该阶段的指标不会持久化到后续fit阶段。可以在tune前移除checkpoint回调，完成后恢复：

# 保存原始回调列表
original_callbacks = trainer.callbacks
# 过滤掉ModelCheckpoint
trainer.callbacks = [cb for cb in original_callbacks if not isinstance(cb, ModelCheckpoint)]
# 执行tune
trainer.tune(model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
# 恢复原始回调
trainer.callbacks = original_callbacks

2. 修改ModelCheckpoint的触发时机 #

将every_n_train_steps=100改为按epoch触发，确保只有在完成完整验证epoch后才会执行checkpoint：

checkpoint_callback = ModelCheckpoint(
    monitor="val_loss",
    every_n_epochs=1,
    verbose=True
)

3. 在tune后手动执行一次验证 #

在tune完成后、fit开始前，手动运行一次验证流程，提前记录val_loss：

trainer.tune(model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
# 手动执行验证
trainer.validate(model, val_dataloaders=val_dataloader)
trainer.fit(model, train_dataloader, val_dataloader)

4. 显式设置log参数 #

在validation_step的log语句中显式指定prog_bar=True和logger=True，确保指标被正确记录：

def validation_step(self, batch, batch_idx):
    loss, embeddings = self._step(batch)
    # 显式设置log参数
    self.log("val_loss", loss, prog_bar=True, logger=True)
    return {
        'val_loss': loss,
        'preds': embeddings
    }

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