【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 WandbLogger

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import random_split, DataLoader

from torchmetrics import Accuracy

from torchvision import transforms
from torchvision.datasets import CIFAR10

import 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_dir
        self.batch_size = batch_size

        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

        self.num_classes = 10

    def 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_samples
        self.val_imgs, self.val_labels = val_samples

    def on_validation_epoch_end(self, trainer, pl_module):
        # 将张量带到 CPU
        val_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 Image
        trainer.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_rate

        self.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 = 1
        input = 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 x

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = 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 loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = 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 loss

    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = 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 loss

    def 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 是小菜一碟。我甚至可以编写自定义回调。