PyTorch Lightning训练报错：传入模型非LightningModule类型

问题背景 #

在Colab GPU环境中使用PyTorch Lightning训练Segformer语义分割模型时，调用trainer.fit()触发类型错误，提示传入的模型不是LightningModule或OptimizedModule类型。

自定义LightningModule代码 #

import torch
import torch.nn as nn
import pytorch_lightning as pl

class LightningSegformerForSemanticSegmentation(pl.LightningModule):
    def __init__(self, segformer):
        super().__init__()
        self.segformer = segformer
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.segformer(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        out = self.segformer(x)
        loss = self.criterion(out, y)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        out = self.segformer(x)
        loss = self.criterion(out, y)
        self.log('val_loss', loss)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

训练代码 #

from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint

early_stop_callback = EarlyStopping(
    monitor="val_loss", 
    min_delta=0.00, 
    patience=10, 
    verbose=False, 
    mode="min",
)

checkpoint_callback = ModelCheckpoint(save_top_k=1, monitor="val_loss")

trainer = pl.Trainer(
    #gpus='1',  
    accelerator='auto',
    callbacks=[early_stop_callback, checkpoint_callback],
    max_epochs=500,
    val_check_interval=len(train_dataloader),
)
trainer.fit(segformer_finetuner)

错误信息 #

INFO:pytorch_lightning.utilities.rank_zero:GPU available: True (cuda), used: True
INFO:pytorch_lightning.utilities.rank_zero:TPU available: False, using: 0 TPU cores
INFO:pytorch_lightning.utilities.rank_zero:IPU available: False, using: 0 IPUs
INFO:pytorch_lightning.utilities.rank_zero:HPU available: False, using: 0 HPUs
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-22-a821d9e5ddff> in <cell line: 19>()
     17 )
     18 #trainer.fit(segformer_finetuner)
---> 19 trainer.fit(segformer_finetuner)

1 frames
/usr/local/lib/python3.9/dist-packages/pytorch_lightning/utilities/compile.py in _maybe_unwrap_optimized(model)
    123     if isinstance(model, pl.LightningModule):
    124         return model
--> 125     raise TypeError(
    126         f"`model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `{type(model).__qualname__}`"
    127     )

TypeError: `model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `SegformerForSemanticSegmentation`

解决方案 #

核心问题 #

传入trainer.fit()的segformer_finetuner是原生SegformerForSemanticSegmentation实例，而非你自定义的LightningSegformerForSemanticSegmentation实例，不符合PyTorch Lightning的要求。

修复步骤 #

1. 正确包装模型：确保用自定义的LightningSegformerForSemanticSegmentation包裹原生Segformer模型
2. 明确传入数据加载器：trainer.fit()需指定训练和验证数据集加载器（若未绑定到模型）

修复后的示例代码：

# 导入原生Segformer模型
from transformers import SegformerForSemanticSegmentation

# 1. 实例化原生Segformer预训练模型
base_segformer = SegformerForSemanticSegmentation.from_pretrained(
    "nvidia/segformer-b0-finetuned-ade-512-512",
    num_labels=你的类别数  # 根据你的任务修改
)

# 2. 用自定义LightningModule包装原生模型
segformer_finetuner = LightningSegformerForSemanticSegmentation(base_segformer)

# 3. 调用trainer.fit时传入数据加载器
trainer.fit(
    model=segformer_finetuner,
    train_dataloaders=train_dataloader,
    val_dataloaders=val_dataloader  # 若有验证集则传入
)

额外检查 #

• 确认自定义的LightningSegformerForSemanticSegmentation类已正确导入到训练代码所在的单元格/文件中
• 若使用Colab，需确保定义模型类的单元格已执行，且segformer_finetuner变量确实是LightningSegformerForSemanticSegmentation类型（可通过print(type(segformer_finetuner))验证）

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