PyTorch Lightning循环训练时模型在epoch 0提前停止原因排查

问题原因 #

你复用了同一个EarlyStopping回调实例，而这个回调是带内部状态的（比如记录最佳验证损失、已等待轮次等）。第一个模型训练完成后，回调的状态没有被重置，后续模型训练时，回调会直接沿用之前的状态，导致刚进入训练就触发停止条件，只跑了epoch0就结束。

修复方案 #

每次训练新模型时，创建新的EarlyStopping回调实例，确保每个模型的回调状态独立。

修改后的核心代码 #

将EarlyStopping的创建移到训练循环内部：

def train_model(train_loader, val_loader, test_loader, num_epochs, device):
    models = [
        BaseModel("efficientnet_b0"),
        BaseModel("mobilenet_v2"),
        BaseModel("mobilenet_v3_large"),
    ]

    for model in models:
        model = model.to(device)
        logger = TensorBoardLogger("lightning_logs", name=model.architecture)
        # 每次循环创建新的EarlyStopping实例
        early_stop_callback = EarlyStopping(
            monitor="val_loss", patience=3, verbose=False, mode="min"
        )
        trainer = pl.Trainer(
            max_epochs=num_epochs,
            callbacks=[early_stop_callback],
            accelerator="gpu",
            devices=1,
            logger=logger,
        )
        trainer.fit(model, train_loader, val_loader)
        trainer.test(model, test_loader)

验证说明 #

修改后，每个模型都会使用独立的回调实例，内部状态互不干扰，所有模型都能按预期完成多轮训练，直到触发EarlyStopping的patience条件或达到max_epochs。

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