使用 PyTorch Lightning 简化训练、验证和测试流程

PyTorch Lightning 能够简化 PyTorch 中模型的训练、验证和测试流程。通过将代码组织到 LightningModule 类中，可以更清晰地定义模型结构和训练逻辑。

以下示例展示了如何使用 PyTorch Lightning 实现一个简单的 MNIST 手写数字识别模型：pythonimport torchfrom torch import nnfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transformsimport pytorch_lightning as pl

class MyModel(pl.LightningModule): def init(self): super().init() self.model = nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10) ) self.loss_fn = nn.CrossEntropyLoss()

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

def training_step(self, batch, batch_idx):        x, y = batch        logits = self.forward(x)        loss = self.loss_fn(logits, y)        self.log('train_loss', loss)        return loss

def validation_step(self, batch, batch_idx):        x, y = batch        logits = self.forward(x)        loss = self.loss_fn(logits, y)        self.log('val_loss', loss)        return loss

def test_step(self, batch, batch_idx):        x, y = batch        logits = self.forward(x)        loss = self.loss_fn(logits, y)        self.log('test_loss', loss)        return loss

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

加载数据transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)val_dataset = datasets.MNIST('data', train=False, download=True, transform=transform)train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)val_loader = DataLoader(val_dataset, batch_size=64)

创建模型和训练器model = MyModel()trainer = pl.Trainer(gpus=1, max_epochs=10)

训练、验证和测试trainer.fit(model, train_loader, val_loader)trainer.test(model, test_dataloaders=val_loader)

代码解析:

1. MyModel 类: 继承自 pl.LightningModule，定义了模型结构和训练逻辑。 - __init__: 初始化模型层和损失函数。 - forward: 定义前向传播过程。 - training_step, validation_step, test_step: 分别定义训练、验证和测试步骤，计算并记录损失。 - configure_optimizers: 定义优化器。2. 数据加载: 使用 torchvision.datasets 加载 MNIST 数据集，并使用 DataLoader 创建数据加载器。3. 模型训练和测试: - 创建 MyModel 实例和 pl.Trainer 实例。 - 调用 trainer.fit 进行训练，传入模型、训练数据加载器和验证数据加载器。 - 调用 trainer.test 进行测试，传入模型和测试数据加载器。

通过使用 PyTorch Lightning，可以更清晰地组织代码，并简化训练、验证和测试流程。