PyTorch Lightning简介
PyTorch Lightning是一个为机器学习研究人员提供的开源、轻量级的Python包装器,它建立在PyTorch之上。
有了这个框架,你不需要记住PyTorch框架的所有微小细节,因为它能处理这些细节。
Pytorch Lightning是一个建立在Pytorch之上的高级框架。正如我们所知,Pytorch已经很不错了。因此,Pytorch Lightning在很多方面更加伟大。
它与Keras和TensorFlow非常相似。该框架的建立是为了使训练神经网络更容易,以及减少所需的训练代码。
该框架可以让你花更少的时间在工程上,更多的时间在研究上。
本教程将介绍这个Pytorch Lightning。我们将实现一个例子来演示如何使用它。
PyTorch Lightning
PyTorch Lightning减轻了负担,让你更专注于研究而不是工程。
这个框架的一个值得注意的特点是,它可以打印警告并给开发者提供机器学习的提示。
许多机器学习开发者对那些倾向于隐藏基础工程的框架持怀疑态度。
这就是为什么众多开发者一开始就喜欢PyTorch而不是TensorFlow。但另一方面,PyTorch Lightning让事情变得更加简单。
因此,在开始使用PyTorch Lightning之前,最好先学习PyTorch的基础知识。
PyTorch Lightning的优势
- 它很容易使用
pip进行安装。 - 该框架的代码倾向于简单、干净、易于复制。这是因为工程代码与主代码是分开的。
- 它支持16位精度。这有助于加快模型训练的速度。
- 它可以运行分布式训练。它支持同时在多台机器上训练。
- 它支持模型检查点。检查点是一种保存实验的当前状态的方法,这样你就可以从你离开的地方继续前进。这有助于恢复以前的状态,以防发生意外,如停电。
- 它很容易与其他流行的机器学习工具集成。例如,它支持谷歌的Tensorboard。
- 与PyTorch相比,它的最小运行速度开销约为300ms,这使得它相当快。
- 它的模型是与硬件无关的。它可以在任何CPU、GPU或TPU机器上运行。
- 如果开发者碰巧在代码中犯了错误,它会打印出警告并给开发者提供机器学习提示。这可能很有帮助。
关于如何使用PyTorch Lightning的演示
安装和导入依赖项
如前所述,Pytorch Lightning是建立在Pytorch之上的。因此,在安装PyTorch Lightning之前,我们仍然需要导入vanilla Pytorch 。
!pip install pytorch-lightning
如果你正在进行Anaconda install ,请使用以下命令。
conda install pytorch-lightning -c conda-forge
安装完成后,我们需要在代码中导入以下依赖项。
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import torch.nn.functional as F
import matplotlib.pyplot as plt
import pytorch_lightning as pl
from pytorch_lightning import Trainer
让我们定义一些超参数,我们将在以后的代码中使用。
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 3
batch_size = 100
learning_rate = 0.001
下面的代码显示了如何描述一个PyTorch Lightning模块。注意这是一个lightning module ,而不是一个torch.nn.Module 。
我们首先需要初始化我们的模型,输入大小为784 神经网络,500 隐藏神经元,以及10 输出类。
由于我们不是在GPU机器上训练,所以我们将epochs的数量设置为只有3 。如果我们增加epochs的数量,将需要很长的时间来训练。
在PyTorch闪电中,forward 定义了预测/推理动作。training_step 定义了训练循环。它是独立于前进的。
我们也在计算交叉熵并返回损失。最后,Tensorboard是机器学习研究者最常用的记录器之一。
为了记录到Tensorboard,你可以使用键log ,它可以从LightningModule的任何方法中调用。
class LitNeuralNet(pl.LightningModule):
def __init__(self, input_size, hidden_size, num_classes):
super(LitNeuralNet, self).__init__()
self.input_size = input_size
self.l1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.l2 = nn.Linear(hidden_size, num_classes)
def forward(self, x):
out = self.l1(x)
out = self.relu(out)
out = self.l2(out)
# no activation and no softmax at the end
return out
def training_step(self, batch, batch_idx):
images, labels = batch
images = images.reshape(-1, 28 * 28)
# Forward pass
outputs = self(images)
loss = F.cross_entropy(outputs, labels)
tensorboard_logs = {'train_loss': loss}
# use key 'log'
return {"loss": loss, 'log': tensorboard_logs}
训练
在训练步骤中,我们将讨论五个函数。
我们有train_dataloader,val_dataloader,validation_step,validation_epoch_end, 和configure_optimizers 函数。
让我们来讨论每个函数的作用。
def train_dataloader(self):
# MNIST dataset
train_dataset = torchvision.datasets.MNIST(
root="./data", train=True, transform=transforms.ToTensor(), download=True
)
# Data loader
train_loader = torch.utils.data.DataLoader(
dataset=train_dataset, batch_size=batch_size, num_workers=2, shuffle=False
)
return train_loader
train_dataloader() 函数生成了训练数据加载器。它允许我们加载我们想在项目中使用的数据集。
在这个项目中,我们正在加载MNIST数据集。工作者的数量,num_workers ,根据你的机器的CPU数量而变化。
大多数计算机有4个CPU。如果是这样,把这个值改为4或你的机器上可用的CPU数量。
def val_dataloader(self):
test_dataset = torchvision.datasets.MNIST(
root="./data", train=False, transform=transforms.ToTensor()
)
test_loader = torch.utils.data.DataLoader(
dataset=test_dataset, batch_size=batch_size, num_workers=2, shuffle=False
)
return test_loader
val_dataloader ,生成验证数据加载器。它允许我们加载验证数据集,让我们对模型性能有一个无偏见的评价。
def validation_step(self, batch, batch_idx):
images, labels = batch
images = images.reshape(-1, 28 * 28)
# Forward pass
outputs = self(images)
loss = F.cross_entropy(outputs, labels)
return {"val_loss": loss}
def validation_epoch_end(self, outputs):
# outputs = list of dictionaries
avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
tensorboard_logs = {'avg_val_loss': avg_loss}
# use key 'log' to load Tensorboard
return {'val_loss': avg_loss, 'log': tensorboard_logs}
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr=learning_rate)
if __name__ == '__main__':
model = LitNeuralNet(input_size, hidden_size, num_classes)
trainer = Trainer(max_epochs=num_epochs)
trainer.fit(model)
在configure_optimizer 函数中,我们可以传入任何你想使用的优化器。对于我们的案例,我们选择了Adam 优化器。
此外,我们可以传入learning_rate 参数;本实验选择的是0.001 的学习率。
在PyTorch Lightning中,任何对这个项目至关重要的东西都会被列出,并以一种在每一个项目中都有凝聚力的方式组织起来。
例如,如果你想知道使用的是哪种数据,你需要到train_dataloader() 函数中去。
如果你想确定使用的是哪个优化器,你只需要看一下configure_optimizers() 函数,就能找到。
你可能已经注意到,这里没有GPU代码或半精度代码。所有这些功能都是在PyTorch Lightning的引擎盖下完成的。
为了更好地理解PyTorch Lightning是如何为我们节省大量时间的,并总结我们在本教程中所学到的知识,让我们来看看它是如何通过删除所有的模板代码来转换一个普通的PyTorch代码。
总结
本教程介绍了PyTorch Lightning,并演示了一个例子来说明如何使用该框架。
PyTorch Lightning是一个强大的库,它可以去除不必要的模板代码,使神经网络的训练更加容易。
你不想把大部分时间花在学习如何在多个GPU之间分配你的模型,或做复杂的工程位。
你更愿意把时间放在研究和开发上,这就是创建这个轻量级python包装器的主要想法。
如果你是PyTorch的常客,你可以在你的下一个项目中试试PyTorch Lightning,或者说,把你的代码转换成PyTorch Lightning,看看它是否达到了预期效果。