【WB 深度学习实验管理】使用 PyTorch Lightning 实现高效的图像分类实验跟踪

本文使用到的 Jupyter Notebook 可在GitHub仓库002文件夹找到，别忘了给仓库点个小心心~~~
https://github.com/LFF8888/FF-Studio-Resources

在机器学习项目中，实验跟踪和结果可视化是至关重要的环节。无论是调整超参数、优化模型架构，还是监控训练过程中的性能变化，清晰的记录和直观的可视化都能显著提升开发效率。然而，许多开发者在实际操作中往往忽视了这一点，导致实验结果难以复现，或者在项目协作中出现混乱。今天，笔者将介绍如何利用 PyTorch Lightning 和 Weights & Biases 这一强大的工具组合，轻松构建和训练一个图像分类模型。通过本文，你将学会如何高效地组织数据管道、定义模型架构，并利用 W&B 实现实验跟踪和结果可视化，让每一次实验都清晰可溯，每一次优化都有据可依。

使用 PyTorch Lightning ⚡️ 进行图像分类

我们将使用 PyTorch Lightning 构建一个图像分类管道。我们将遵循这个风格指南来提高代码的可读性和可重复性。这里有一个很酷的解释：使用 PyTorch Lightning 进行图像分类。

设置 PyTorch Lightning 和 W&B

对于本教程，我们需要 PyTorch Lightning（这不是很明显吗！）和 Weights and Biases。

!pip install lightning torchvision -q
# 安装 weights and biases
!pip install wandb -qU

你需要这些导入。

import lightning.pytorch as pl
# 你最喜欢的机器学习跟踪工具
from lightning.pytorch.loggers import WandbLoggerimport torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import random_split, DataLoaderfrom torchmetrics import Accuracyfrom torchvision import transforms
from torchvision.datasets import CIFAR10import wandb

现在你需要登录到你的 wandb 账户。

wandb.login()

🔧 DataModule - 我们应得的数据管道

DataModules 是一种将数据相关的钩子与 LightningModule 解耦的方式，以便你可以开发与数据集无关的模型。
它将数据管道组织成一个可共享和可重用的类。一个 datamodule 封装了 PyTorch 中数据处理的五个步骤：

下载 / 分词 / 处理。
清理并（可能）保存到磁盘。
加载到 Dataset 中。
应用转换（旋转、分词等）。
包装到 DataLoader 中。

了解更多关于 datamodules 的信息这里。让我们为 Cifar-10 数据集构建一个 datamodule。

class CIFAR10DataModule(pl.LightningDataModule):def __init__(self, batch_size, data_dir: str = './'):super().__init__()self.data_dir = data_dirself.batch_size = batch_sizeself.transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])self.num_classes = 10def prepare_data(self):CIFAR10(self.data_dir, train=True, download=True)CIFAR10(self.data_dir, train=False, download=True)def setup(self, stage=None):# 为 dataloaders 分配训练/验证数据集if stage == 'fit' or stage is None:cifar_full = CIFAR10(self.data_dir, train=True, transform=self.transform)self.cifar_train, self.cifar_val = random_split(cifar_full, [45000, 5000])# 为 dataloader(s) 分配测试数据集if stage == 'test' or stage is None:self.cifar_test = CIFAR10(self.data_dir, train=False, transform=self.transform)def train_dataloader(self):return DataLoader(self.cifar_train, batch_size=self.batch_size, shuffle=True)def val_dataloader(self):return DataLoader(self.cifar_val, batch_size=self.batch_size)def test_dataloader(self):return DataLoader(self.cifar_test, batch_size=self.batch_size)

📱 Callbacks

回调是一个独立的程序，可以在项目之间重用。PyTorch Lightning 提供了一些内置回调，这些回调经常被使用。
了解更多关于 PyTorch Lightning 中的回调这里。

内置回调

在本教程中，我们将使用 Early Stopping 和 Model Checkpoint 内置回调。它们可以传递给 Trainer。

自定义回调

如果你熟悉自定义 Keras 回调，那么在 PyTorch 管道中实现相同功能的能力只是锦上添花。
由于我们正在进行图像分类，能够可视化模型对一些样本图像的预测可能很有帮助。这种形式的回调可以帮助在早期阶段调试模型。

class ImagePredictionLogger(pl.callbacks.Callback):def __init__(self, val_samples, num_samples=32):super().__init__()self.num_samples = num_samplesself.val_imgs, self.val_labels = val_samplesdef on_validation_epoch_end(self, trainer, pl_module):# 将张量带到 CPUval_imgs = self.val_imgs.to(device=pl_module.device)val_labels = self.val_labels.to(device=pl_module.device)# 获取模型预测logits = pl_module(val_imgs)preds = torch.argmax(logits, -1)# 将图像记录为 wandb Imagetrainer.logger.experiment.log({"examples":[wandb.Image(x, caption=f"Pred:{pred}, Label:{y}")for x, pred, y in zip(val_imgs[:self.num_samples],preds[:self.num_samples],val_labels[:self.num_samples])]})

🎺 LightningModule - 定义系统

LightningModule 定义了一个系统，而不是一个模型。在这里，系统将所有研究代码分组到一个类中，使其自包含。LightningModule 将你的 PyTorch 代码组织成 5 个部分：

计算 (__init__)。
训练循环 (training_step)
验证循环 (validation_step)
测试循环 (test_step)
优化器 (configure_optimizers)

因此，可以构建一个与数据集无关的模型，并且可以轻松共享。让我们为 Cifar-10 分类构建一个系统。

class LitModel(pl.LightningModule):def __init__(self, input_shape, num_classes, learning_rate=2e-4):super().__init__()# 记录超参数self.save_hyperparameters()self.learning_rate = learning_rateself.conv1 = nn.Conv2d(3, 32, 3, 1)self.conv2 = nn.Conv2d(32, 32, 3, 1)self.conv3 = nn.Conv2d(32, 64, 3, 1)self.conv4 = nn.Conv2d(64, 64, 3, 1)self.pool1 = torch.nn.MaxPool2d(2)self.pool2 = torch.nn.MaxPool2d(2)n_sizes = self._get_conv_output(input_shape)self.fc1 = nn.Linear(n_sizes, 512)self.fc2 = nn.Linear(512, 128)self.fc3 = nn.Linear(128, num_classes)self.accuracy = Accuracy(task="multiclass", num_classes=num_classes)# 返回从卷积块进入线性层的输出张量的大小。def _get_conv_output(self, shape):batch_size = 1input = torch.autograd.Variable(torch.rand(batch_size, *shape))output_feat = self._forward_features(input)n_size = output_feat.data.view(batch_size, -1).size(1)return n_size# 返回卷积块的特征张量def _forward_features(self, x):x = F.relu(self.conv1(x))x = self.pool1(F.relu(self.conv2(x)))x = F.relu(self.conv3(x))x = self.pool2(F.relu(self.conv4(x)))return x# 将在推理期间使用def forward(self, x):x = self._forward_features(x)x = x.view(x.size(0), -1)x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = F.log_softmax(self.fc3(x), dim=1)return xdef training_step(self, batch, batch_idx):x, y = batchlogits = self(x)loss = F.nll_loss(logits, y)# 训练指标preds = torch.argmax(logits, dim=1)acc = self.accuracy(preds, y)self.log('train_loss', loss, on_step=True, on_epoch=True, logger=True)self.log('train_acc', acc, on_step=True, on_epoch=True, logger=True)return lossdef validation_step(self, batch, batch_idx):x, y = batchlogits = self(x)loss = F.nll_loss(logits, y)# 验证指标preds = torch.argmax(logits, dim=1)acc = self.accuracy(preds, y)self.log('val_loss', loss, prog_bar=True)self.log('val_acc', acc, prog_bar=True)return lossdef test_step(self, batch, batch_idx):x, y = batchlogits = self(x)loss = F.nll_loss(logits, y)# 验证指标preds = torch.argmax(logits, dim=1)acc = self.accuracy(preds, y)self.log('test_loss', loss, prog_bar=True)self.log('test_acc', acc, prog_bar=True)return lossdef configure_optimizers(self):optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)return optimizer

🚋 训练和评估

现在我们已经使用 DataModule 组织了数据管道，并使用 LightningModule 组织了模型架构和训练循环，PyTorch Lightning Trainer 为我们自动化了其他所有内容。

Trainer 自动化了以下内容：

Epoch 和 batch 迭代
调用 optimizer.step()、backward、zero_grad()
调用 .eval()，启用/禁用梯度
保存和加载权重
Weights and Biases 日志记录
多 GPU 训练支持
TPU 支持
16 位训练支持

dm = CIFAR10DataModule(batch_size=32)
# 要访问 x_dataloader，我们需要调用 prepare_data 和 setup。
dm.prepare_data()
dm.setup()# 自定义 ImagePredictionLogger 回调所需的样本，用于记录图像预测。
val_samples = next(iter(dm.val_dataloader()))
val_imgs, val_labels = val_samples[0], val_samples[1]
val_imgs.shape, val_labels.shape

model = LitModel((3, 32, 32), dm.num_classes)# 初始化 wandb logger
wandb_logger = WandbLogger(project='wandb-lightning', job_type='train')# 初始化 Callbacks
early_stop_callback = pl.callbacks.EarlyStopping(monitor="val_loss")
checkpoint_callback = pl.callbacks.ModelCheckpoint()# 初始化一个 trainer
trainer = pl.Trainer(max_epochs=2,logger=wandb_logger,callbacks=[early_stop_callback,ImagePredictionLogger(val_samples),checkpoint_callback],)# 训练模型 ⚡🚅⚡
trainer.fit(model, dm)# 在保留的测试集上评估模型 ⚡⚡
trainer.test(dataloaders=dm.test_dataloader())# 关闭 wandb run
wandb.finish()

最终想法

我来自 TensorFlow/Keras 生态系统，发现 PyTorch 虽然是一个优雅的框架，但有点让人不知所措。这只是我的个人经验。在探索 PyTorch Lightning 时，我意识到几乎所有让我远离 PyTorch 的原因都得到了解决。以下是我兴奋的快速总结：

过去：传统的 PyTorch 模型定义通常分散在各个地方。模型在某个 model.py 脚本中，训练循环在 train.py 文件中。需要来回查看才能理解管道。
现在：LightningModule 作为一个系统，模型定义与 training_step、validation_step 等一起定义。现在它是模块化的且可共享的。
过去：TensorFlow/Keras 最棒的部分是输入数据管道。他们的数据集目录丰富且不断增长。PyTorch 的数据管道曾经是最大的痛点。在普通的 PyTorch 代码中，数据下载/清理/准备通常分散在许多文件中。
现在：DataModule 将数据管道组织成一个可共享和可重用的类。它只是 train_dataloader、val_dataloader(s)、test_dataloader(s) 以及匹配的转换和数据处理/下载步骤的集合。
过去：使用 Keras，可以调用 model.fit 来训练模型，调用 model.predict 来运行推理。model.evaluate 提供了一个简单而有效的测试数据评估。这在 PyTorch 中不是这样。通常会找到单独的 train.py 和 test.py 文件。
现在：有了 LightningModule，Trainer 自动化了一切。只需调用 trainer.fit 和 trainer.test 来训练和评估模型。
过去：TensorFlow 喜欢 TPU，PyTorch…嗯！
现在：使用 PyTorch Lightning，可以轻松地在多个 GPU 上训练相同的模型，甚至在 TPU 上。哇！
过去：我是回调的忠实粉丝，更喜欢编写自定义回调。像 Early Stopping 这样简单的事情曾经是传统 PyTorch 的讨论点。
现在：使用 PyTorch Lightning，使用 Early Stopping 和 Model Checkpointing 是小菜一碟。我甚至可以编写自定义回调。