PyTorch Lightning基础教程：快速构建和训练模型-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00167/article/details/148360259

PyTorch Lightning基础教程：快速构建和训练模型

pytorch-lightning Lightning-AI/pytorch-lightning: PyTorch Lightning 是一个轻量级的高级接口，用于简化 PyTorch 中深度学习模型的训练流程。它抽象出了繁杂的工程细节，使研究者能够专注于模型本身的逻辑和实验设计，同时仍能充分利用PyTorch底层的灵活性。项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning

前言

在深度学习项目中，编写训练循环往往是重复且容易出错的部分。PyTorch Lightning框架通过提供高级抽象，让开发者能够专注于模型设计而非工程细节。本文将介绍如何使用PyTorch Lightning快速构建和训练一个简单的自编码器模型。

环境准备

首先确保已安装必要的库：

PyTorch
PyTorch Lightning
TorchVision

模型架构设计

我们将构建一个简单的自编码器，包含编码器和解码器两部分：

import torch
from torch import nn

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Sequential(
            nn.Linear(28 * 28, 64),  # 将784维输入压缩到64维
            nn.ReLU(),
            nn.Linear(64, 3)         # 进一步压缩到3维潜在空间
        )

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

class Decoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Sequential(
            nn.Linear(3, 64),       # 从3维潜在空间扩展
            nn.ReLU(),
            nn.Linear(64, 28 * 28)   # 重建原始784维输出
        )

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

这个自编码器结构简单但完整，展示了特征压缩和重建的基本原理。

创建LightningModule

PyTorch Lightning的核心是LightningModule，它封装了模型定义、训练逻辑和优化器配置：

import lightning as L
import torch.nn.functional as F

class LitAutoEncoder(L.LightningModule):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def training_step(self, batch, batch_idx):
        x, _ = batch  # 忽略标签
        x = x.view(x.size(0), -1)  # 展平图像
        z = self.encoder(x)  # 编码
        x_hat = self.decoder(z)  # 解码
        loss = F.mse_loss(x_hat, x)  # 计算重建损失
        return loss

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

关键点说明：

training_step定义了前向传播和损失计算
configure_optimizers返回优化器实例
无需手动编写反向传播代码

准备数据

使用MNIST数据集作为示例：

from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import DataLoader

transform = transforms.ToTensor()
dataset = MNIST(os.getcwd(), download=True, transform=transform)
train_loader = DataLoader(dataset)

模型训练

PyTorch Lightning的Trainer类封装了完整的训练流程：

# 初始化模型
autoencoder = LitAutoEncoder(Encoder(), Decoder())

# 创建训练器并开始训练
trainer = L.Trainer()
trainer.fit(model=autoencoder, train_dataloaders=train_loader)

Trainer自动处理了以下细节：

训练循环的迭代
梯度计算和参数更新
训练日志记录
硬件加速(如GPU)的配置

与传统PyTorch代码对比

传统PyTorch需要手动编写训练循环：

autoencoder = LitAutoEncoder(Encoder(), Decoder())
optimizer = autoencoder.configure_optimizers()

for batch_idx, batch in enumerate(train_loader):
    loss = autoencoder.training_step(batch, batch_idx)
    
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

而PyTorch Lightning的优势在于：