PyTorch Lightning 是一个轻量级 PyTorch 包装器，它为训练 PyTorch 模型提供了高级接口。它旨在简化和标准化训练循环，从而更轻松地为深度学习项目编写更简洁、更模块化的代码。PyTorch Lightning 引入了一组抽象和约定，这些抽象和约定消除了样板代码，使研究人员和从业者能够更多地关注模型架构和实验配置。

2、PyTorch存在的问题

虽然PyTorch提供动态计算图，与静态计算图框架相比，具有更大的灵活性和易用性。但是使用PyTorch依然存在问题，其问题在于存在大量的样板代码，也就是说，每次启动项目时，都必须重写这些训练和测试循环代码。可以说，训练循环就是 Lightning 从 PyTorch 中脱颖而出的原因。在 PyTorch 中，我们定义完整的训练循环，而在 Lightning 中，我们使用 Trainer()来完成这项工作。

在PyTorch中，我们的训练循环是这么写的：

epochs = 5

for i in range(epochs):
    train_loss = 0.0
    for data,label in trainloader:
        if is_gpu:
            data, label = data.cuda(), label.cuda()
        output = model(data)
        optimizer.zero_grad()
        loss = criterion(output,label)
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * data.size(0)

而在PyTorch-Lightning中，我们可以这么写：

def training_step(self, train_batch, batch_idx):
    x, y = train_batch
    logits = self.forward(x)
    loss = self.loss(logits,y)
    return loss

当然，仅仅看到这里，也许大家还未发现PyTorch Lightning的魅力，请继续阅读下文吧。

3、从 Pytorch 到 Pytorch Lightning

让我们比较一下使用 PyTorch 和 PyTorch Lightning 在 MNIST 数据集上简单 3 层神经网络的训练和验证循环。关键要素包括模型、数据集（MNIST）、优化器和损失函数。

3.1、基本PyTorch的神经网络

3.1.1、导入必要的库和模块

该代码首先导入用于构建和训练神经网络的必要库和模块。这些包括 torch、torch.nn、torch.optim 和 torchvision。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

3.1.2、定义神经网络模型

接下来，该代码使用 PyTorch 的 nn.Module 类。该模型由三个全连接层（fc1、fc2 和 fc3）组成，分别具有 256、128 和 10 个神经元。输出层有 10 个神经元，对应于 MNIST 数据集中的 10 类数字。

forward 方法定义神经网络的 forward 传递，其中输入通过每一层并使用 ReLU 激活函数进行转换。

# Define the model
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

3.1.3、加载数据集

然后，该代码使用 torchvision.datasets.MNIST 加载 MNIST 数据集。使用转换对数据集进行预处理。Compose 方法，该方法对数据应用一系列转换。在这种情况下，数据被转换为张量并标准化为平均值 0.1307 和标准差 0.3081。

创建 train_dataset 和 test_dataset 对象并将其加载到 train_loader 和 test_loader 中，它们是处理数据批处理和随机排序的 PyTorch DataLoader 对象。

# Load the dataset
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST("data", train=True, download=True, transform=transform)
test_dataset = datasets.MNIST("data", train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)

3.1.4、初始化模型、优化器和损失函数

初始化神经网络模型，并定义随机梯度下降（SGD）优化器和交叉熵损失函数。

# Initialize the model
model = NeuralNetwork()

# Define the optimizer and loss function
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

3.1.5、训练和验证模型

该代码定义了两个函数 train 和 validate，它们处理神经网络的训练和验证。

train 函数接收模型、训练数据加载器、优化器和损失函数，并在数据上批量训练模型。梯度将累积，并使用 SGD 优化器更新模型权重。 validate 函数接收模型、测试数据加载器和损失函数，并根据测试数据评估模型。计算测试损失，并将模型的准确性计算为正确预测的百分比。最后，代码使用 train 和 validate 函数训练和验证神经网络 10 个 epoch。

# Training loop
def train(model, train_loader, optimizer, criterion):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# Validation loop
def validate(model, test_loader, criterion):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print("Validation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)".format(
        test_loss, correct, len(test_loader.dataset), accuracy))

# Train and validate the model
for epoch in range(10):
    train(model, train_loader, optimizer, criterion)
    validate(model, test_loader, criterion)

输出：

Validation set: Average loss: 0.0052, Accuracy: 9050/10000 (90.50%)
Validation set: Average loss: 0.0041, Accuracy: 9267/10000 (92.67%)
Validation set: Average loss: 0.0034, Accuracy: 9369/10000 (93.69%)
Validation set: Average loss: 0.0030, Accuracy: 9445/10000 (94.45%)
Validation set: Average loss: 0.0026, Accuracy: 9506/10000 (95.06%)
Validation set: Average loss: 0.0024, Accuracy: 9555/10000 (95.55%)
Validation set: Average loss: 0.0021, Accuracy: 9597/10000 (95.97%)
Validation set: Average loss: 0.0020, Accuracy: 9641/10000 (96.41%)
Validation set: Average loss: 0.0018, Accuracy: 9663/10000 (96.63%)
Validation set: Average loss: 0.0017, Accuracy: 9680/10000 (96.80%)

3.2、PyTorch 闪电

3.2.1、首先，进行必要的导入，包括 PyTorch、PyTorch Lightning、MNIST 数据集和 Adam 优化器。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torch.optim import Adam
import pytorch_lightning as pl

3.2.2、定义神经网络模型

接下来，定义 MyModel 类，该类继承自 pl。LightningModule 的 LightningModule 中。此类定义神经网络架构、前向传递、训练步骤、验证步骤和优化器的配置。

在 init 方法中，神经网络架构是使用 PyTorch 的 nn.顺序模块。它由三个具有 ReLU 激活函数的全连接层和一个用于输出概率的最终 softmax 层组成。

forward 方法接受输入张量 x，将其重塑为具有正确的维数，然后使用 self.model 模块将其传递给神经网络。

class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        x = x.view(x.size(0), -1)  # Reshape the input
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.CrossEntropyLoss()(y_hat, y)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.CrossEntropyLoss()(y_hat, y)
        self.log("val_loss", loss)
        # Calculate accuracy
        correct = (y_hat.argmax(1) == y).sum().item()
        total = y.size(0)
        self.log("accuracy", correct / total, on_step=False, on_epoch=True, prog_bar=True)

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

3.2.3、加载数据集

MNIST 数据集是使用 torchvision.datasets 中的 MNIST 类加载的。训练集和验证集被拆分为单独的 DataLoader 对象，用于训练和验证。

# Load the MNIST dataset
train_dataset = MNIST(root=".", train=True, transform=ToTensor(), download=True)
val_dataset = MNIST(root=".", train=False, transform=ToTensor())

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=64)

3.2.4、初始化模型

# Initialize the model
model = MyModel()

3.2.5、训练模型

trainer 对象是使用 pl 创建的。教练。max_epochs 参数设置为 10，这意味着模型将训练 10 个 epoch。如果 GPU 不可用，则 accelerator 参数设置为 “cpu”，否则设置为 “cpu”。

# Initialize the trainer
trainer = pl.Trainer(max_epochs=10, accelerator = "cpu" if torch.cuda.is_available() else "cpu")

3.2.6、模型拟合

最后，使用 trainer.fit 方法训练模型。它接受 model 对象以及 train_loader 和 val_loader 对象作为参数。

# Train the model
trainer.fit(model, train_loader, val_loader)

3.2.7、验证模型

trainer.validate(model, val_loader)

输出：

[{"val_loss": 1.4881926774978638, "accuracy": 0.9729999899864197}]

4、在代码风格方面，Pytorch与Pytorch Lightning不同之处

特性	PyTorch	PyTorch Lightning
训练循环	显式编写训练过程的for循环，包括处理数据到GPU、计算梯度和更新模型参数等方面。	使用trainer.fit更简洁。用户定义 hook 和 callback 来自定义行为，而无需直接修改训练循环。
优化器和调度器	在训练循环中定义和配置	在方法configure_optimizers中定义，使用专用方法进行标准化
日志记录指标	记录指标和跟踪实验需要使用TensorBoard等工具或自定义记录器手动实现。	内置支持各种日志记录框架（TensorBoard、CSV等）和实验跟踪平台（例如WandB、Comet）
GPU和分布式训练	需要手动将模型和数据显式移动到GPU，管理分布式训练，并处理多GPU场景。	根据用户配置自动。用户可以指定 GPU 的数量。

_{本教程共117节，当前为第5节!}
_{本教程最新修订时间为：2026-05-08 11:10:53}

网站公告：2026年7月招聘兼职>>>>>>