PyTorch Lightning的举例介绍

996 阅读6分钟

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,看看它是否达到了预期效果。