PyTorch Lightning（pl），一个轻量级的深度学习框架

最近刚做完毕设，发现身边的人基本都在用PyTorch或者tensorflow作为深度学习的框架，接触到pl的人少之又少。在深度学习项目中，代码的复用性和训练的调试一直是让人头疼的问题。每次进行新的实验或调优模型时，我们常常需要花费大量时间来调整代码，解决各种意想不到的bug，导致进展缓慢。有没有一种方法可以简化这一过程，让我们更专注于模型的设计和实验结果呢？

答案是肯定的——这就是 PyTorch Lightning。PyTorch Lightning 是一个轻量级的框架，它基于 PyTorch，为我们提供了一种简洁高效的方式来组织代码、管理训练过程，并提升代码的可读性和复用性。通过使用 PyTorch Lightning，我们可以更专注于核心研究问题，而不必为繁琐的代码细节所困扰。

然而在GitHub以及各种计算机领域的社区，极少提及到这个框架的相关概念，关于这一点在新手的学习中也是较为困难，唯一的途径就是查看官方文档，链接我贴在下面。写这篇教程就是想帮助对这个框架感兴趣的朋友来入门。PyTorch Lightning 官方文档

在这篇教程中，我们将深入探讨 PyTorch Lightning 的基本概念、使用方法，以及如何在实际项目中应用它来提升工作效率。无论你是深度学习的新手，还是经验丰富的研究者，相信这篇教程都会给你带来启发和帮助。

---

对于未使用过该框架的用户，需要下载，打开终端在自己的虚拟环境中输入以下命令进行安装：

pip install pytorch-lightning

在 Python 解释器中输入以下代码，确保安装成功：

import pytorch_lightning as pl
print(pl.__version__)

---

对于传统的深度学习而言，我们就拿最简单的回归或者分类任务举例，完成一个深度学习的项目最少需要以下步骤：

1、准备数据集，并对数据进行预处理，数据清洗等

2、创建Dataset实例和DataLoader（数据加载器），确保模型在训练的时候能够分批次将数据传入

3、训练过程，包括优化器的定义，正向传播反向传播，计算损失梯度更新等……

4、测试

这些流程涉及到很多细微的操作，比如需要使用到各种回调（早停），保存模型参数，如何使用gpu并保证数据能够加载到gpu上面等等，使用传统的PyTorch都需要自己定义而且有很多繁琐的工作。

PyTorch Lightning 自动处理了许多繁琐的操作。例如，使用 PyTorch Lightning，可以轻松实现早停回调和模型参数保存，而无需编写复杂的代码。只需简单地配置回调函数，Lightning 就会自动监控验证损失并在合适的时候停止训练，并且会在训练过程中自动保存和加载最佳模型参数。此外，Lightning 提供了简洁的接口来使用 GPU，只需简单地配置设备参数，Lightning 就会自动处理数据和模型的加载过程。

那么我们就从上述步骤开始，编写lightning模块的代码。

---

PyTorch Lightning 教程：MNIST 分类任务

本文将带你一步步使用 PyTorch Lightning 完成一个 MNIST 分类任务。我们将从导入必要的模块开始，并详细解释每个步骤的具体操作。

导入模块

首先，我们需要导入一些必要的模块，包括 PyTorch、PyTorch Lightning 以及其他辅助库。

import os
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
from torchvision.datasets import MNIST

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import TensorBoardLogger

DataModule从数据集准备开始

对于PyTorch通常是重写Dataset类，我们这里也可以这么写，但是pl的DataModule有为你提供准备数据集的部分，因此咱们直接上手吧。

class MnistDataModule(pl.LightningDataModule): # 用于加载数据模块
    def __init__(self, data_dir, batch_size, num_workers): # 这里定义一些超参数并初始化
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers

    def prepare_data(self): # 重写prepare_data方法，需要读入数据集到内存的时候运行
        datasets.MNIST(self.data_dir, train=True, download=True)
        datasets.MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage): # 准备数据集并划分成训练，验证，测试
        entire_dataset = datasets.MNIST(
            root=self.data_dir,
            train=True,
            transform=transforms.ToTensor(),
            download=False,
        )
        self.train_ds, self.val_ds = random_split(entire_dataset, [50000, 10000])
        self.test_ds = datasets.MNIST(
            root=self.data_dir,
            train=False,
            transform=transforms.ToTensor(),
            download=False,
        )

# 以下创建训练验证测试的数据加载器
    def train_dataloader(self): 
        return DataLoader(
            self.train_ds,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=True,
        )

    def val_dataloader(self):
        return DataLoader(
            self.val_ds,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=False,
        )

    def test_dataloader(self):
        return DataLoader(
            self.test_ds,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            shuffle=False,
        )

这个DataModule模块的重写可以直接封装好数据集和数据加载器的准备，也就相当于我们实例化模型的时候可以直接把这个module给到模型，模型直接调forward进行前向传播。

LightningModule 搭建模型

继承pl.LightningModule并且定义三个最重要的生命周期方法

• training_step
• validation_step
• test_step

pl.LightningModule可以视为torch.nn.Module的再包装

class NN(pl.LightningModule):
  def __init__(self, input_size, learning_rate, num_classes):
    super().__init__()
    self.lr = learning_rate # 学习率
    self.fc1 = nn.Linear(input_size, 50)
    self.fc2 = nn.Linear(50, num_classes)
    self.loss_fn = nn.CrossEntropyLoss() # 定义损失函数
    self.accuracy = torchmetrics.Accuracy(
        task="multiclass", num_classes=num_classes
    ) # accuracy
    self.f1_score = torchmetrics.F1Score(
        task="multiclass", num_classes=num_classes
    ) # F1 Score

  def forward(self, x): # 这里定义前向传播方法，神经网络层
    x = F.relu(self.fc1(x))
    x = self.fc2(x)
    return x

# 从这里开始定义模型的训练，验证，测试过程
# 训练，定义每一个step的loss如何计算，也包括其他指标
# 拿到的数据来自train_dataloader
  def training_step(self, batch, batch_idx):
    loss, scores, y = self._common_step(batch, batch_idx) # 这里调用了forward来训练
    accuracy = self.accuracy(scores, y)
    f1_score = self.f1_score(scores, y)
    self.log_dict(
        {
            "train_loss": loss,
            "train_accuracy": accuracy,
            "train_f1_score": f1_score,
          },
          on_step=False,
          on_epoch=True,
          prog_bar=True,

        ) # 保存训练指标
    return {"loss": loss, "scores": scores, "y": y}
    
# 验证，定义验证step，计算loss
# 拿到的数据来自val_dataloader
  def validation_step(self, batch, batch_idx):
        loss, scores, y = self._common_step(batch, batch_idx)
        self.log("val_loss", loss)
        return loss

# 测试，定义测试指标
# 评估最终模型性能，数据来自test_dataloader    
  def test_step(self, batch, batch_idx):
        loss, scores, y = self._common_step(batch, batch_idx)
        self.log("test_loss", loss)
        return loss

# 这里定义了如何从batch中拿到资料
#  也可以直接在上面的地方定义
  def _common_step(self, batch, batch_idx):
        x, y = batch
        x = x.reshape(x.size(0), -1)
        scores = self.forward(x) # 通过神经网络拟合结果
        loss = self.loss_fn(scores, y)
        return loss, scores, y

# 预测的流程，可以不写
# 实例化trainer之后可以调用trainer.predict()来预测    
  def predict_step(self, batch, batch_idx):
        x, y = batch
        x = x.reshape(x.size(0), -1)
        scores = self.forward(x)
        preds = torch.argmax(scores, dim=1)
        return 

# 优化器
  def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr)

这个模块的内容有点多，但是看到这里并且能看懂的话已经完成大部分工作了。总体而言，就是定义了model然后将训练验证测试的流程封装在这个模块中，我们可以直接在后续的Trainer中实例化这个模块并调用其中的方法。

Trainer

Trainer抽象了整个training loop，而且有很多开箱即用的callbacks，我们只需要调用其中的方法，不需要自己定义。我自己常用的有EarlyStopping、ModelCheckpoint，基本上将这些组织起来就是一个非常完整的training过程。

train.fit()会在不同阶段搭配DataModule去调用training_step和validation_step

我们先定义所有超参数：

import os

# Training hyperparameters
INPUT_SIZE = 784
NUM_CLASSES = 10
LEARNING_RATE = 0.001
BATCH_SIZE = 64
NUM_EPOCHS = 3

# Dataset
DATA_DIR = "dataset/"
NUM_WORKERS = os.cpu_count()

# Compute related
ACCELERATOR = "auto"
DEVICES = "auto"
PRECISION = 16

然后实例化该模型，具体如下：

# 实例化model，传入超参数
model = NN(
    input_size=INPUT_SIZE,
    learning_rate=LEARNING_RATE,
    num_classes=NUM_CLASSES,
) 

# 实例化数据模块
dm = MnistDataModule(
    data_dir=DATA_DIR,
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
)

# 在这里使用Trainer
# 我们只需要定义accelerator参数，pl就可以自动将数据和模型上gpu
trainer = pl.Trainer(
    devices=ACCELERATOR,
    accelerator=DEVICES, 
    min_epochs=1,
    max_epochs=3,
    precision=PRECISION,
)

# 训练
trainer.fit(model, dm)

# 验证
trainer.validate(model, dm)

# 测试
trainer.test(model, dm)

PyTorch Lightning拥有非常好的UI，我们在训练的过程中，可以清楚地看到训练的各种指标，包括参数量，训练流程，训练的loss和其他指标等。

这里的代码我们并没有使用到EarlyStopping 和ModelCheckPoint回调，下面是如何使用这两者来回调的代码：

# 导入
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint

# 定义早停实例
# 以val_loss为指标，越小越好
# pl会在训练的过程中通过val_loss来更新
earlystopping = EarlyStopping('val_loss', patience=5, verbose=True, 
                               min_delta=0.00, mode="min")

# 保存最佳模型
# 以val_loss 为指标，越小越好
# pl会根据指标自动保存权重并更新最佳的模型
checkpoint_callback = ModelCheckpoint(monitor='val_loss', # 指标
                                        mode='min', # 根据大还是小来更新
                                        dirpath='./', # 保存路径
                                        filename='{epoch:02d}-{val_loss:.2f}', # 保存文件名
                                        save_top_k=1 # 保存最佳的模型个数
                                        )
callback = [earlystopping, checkpoint_callback] # 封装回调函数

# 实例化Trainer
trainer = pl.Trainer(
    devices=ACCELERATOR,
    accelerator=DEVICES, 
    min_epochs=1,
    max_epochs=3,
    precision=PRECISION,
    callbacks=callback, # 只需要在这里加入callback参数即可
)
# 后面的训练过程一样，省略

就此为止，我们完成了整个深度学习的训练过程，我们展示了如何使用 PyTorch Lightning 来简化深度学习项目中的各类繁琐操作。从导入模块、定义模型、准备数据到训练和测试模型，PyTorch Lightning 提供了一种简洁高效的方式来组织代码和管理训练过程。

此外，PyTorch Lightning 还提供了丰富的 API 和功能，例如支持分布式训练、多 GPU 训练、混合精度训练、丰富的回调函数（如学习率调度、日志记录、模型监控等），以及与各种日志记录工具（如 TensorBoard、MLflow 等）的集成。这些功能使得 PyTorch Lightning 不仅适用于简单的实验项目，也能胜任大规模生产级别的深度学习任务。

具体的 API 介绍和更多高级功能，请参阅 PyTorch Lightning 的官方文档。希望这个教程对你有所帮助，能够提升你的深度学习开发效率。