最近刚做完毕设,发现身边的人基本都在用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 的官方文档。希望这个教程对你有所帮助,能够提升你的深度学习开发效率。