如何将原生PyTorch格式模型无需复制粘贴导入LightningModule

你可以通过两种无代码复制的方案直接复用现有模型的__init__、forward逻辑以及全部结构，无需修改原有模型的任何代码：

方案1：组合封装（推荐，通用性更强） #

直接把原生模型作为参数传入Lightning模块，训练、推理逻辑完全独立封装，可适配任意原生PyTorch模型：

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

# 原有原生PyTorch模型完全不需要修改
class NormalAutoEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
        self.decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

    def forward(self, x):
        embedding = self.encoder(x)
        return embedding

class LitAutoEncoder(pl.LightningModule):
    def __init__(self, raw_model: nn.Module, lr: float = 1e-3):
        super().__init__()
        # 直接注入原生模型实例，复用全部结构与方法
        self.model = raw_model
        self.lr = lr
        self.loss_fn = nn.MSELoss()

    # 直接复用原生模型的forward逻辑，推理行为和原有完全一致
    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, _ = batch
        x = x.view(x.size(0), -1)
        # 可直接访问原生模型的所有属性
        z = self.model.encoder(x)
        x_hat = self.model.decoder(z)
        loss = self.loss_fn(x_hat, x)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr)

使用方式：

# 初始化原生模型
raw_ae = NormalAutoEncoder()
# 封装为Lightning模块
lit_ae = LitAutoEncoder(raw_model=raw_ae)
# 后续直接使用Lightning Trainer训练、推理即可

方案2：多继承实现（更简洁） #

如果希望直接在Lightning模块下访问原生模型的所有属性，不需要额外加层级，可以用多继承的方式：

class LitAutoEncoder(pl.LightningModule, NormalAutoEncoder):
    def __init__(self, lr: float = 1e-3):
        # 分别初始化两个父类
        pl.LightningModule.__init__(self)
        NormalAutoEncoder.__init__(self)
        self.lr = lr
        self.loss_fn = nn.MSELoss()

    # forward方法自动继承NormalAutoEncoder的实现，无需重写
    def training_step(self, batch, batch_idx):
        x, _ = batch
        x = x.view(x.size(0), -1)
        # 直接访问原有模型的属性，不需要加.model前缀
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = self.loss_fn(x_hat, x)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr)

两种方案都完全不需要复制原有模型的结构代码，实现了逻辑的完全复用。

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