如何从Pytorch 到 Pytorch Lightning | 简要介绍

Posted Tina姐

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了如何从Pytorch 到 Pytorch Lightning | 简要介绍相关的知识,希望对你有一定的参考价值。

这篇文章主要介绍为什么使用pytorch时,需要使用Lightning的最常见问题。由Pytorch Lightning的主创团队编写(William Falcon),经本文翻译。

PyTorch非常易于使用,可以构建复杂的AI模型。但是一旦研究变得复杂,并且将诸如多GPU训练,16位精度和TPU训练之类的东西混在一起,用户很可能引入Bug。

PyTorch Lightning完全解决了这个问题。Lightning会构建您的PyTorch代码,以便可以抽象出训练的细节。这使得AI研究可扩展且可快速迭代

该系列上一篇地址: pytorch-lightning入门(一)—— 初了解

1 谁适用PyTorch Lightning?

PyTorch Lightning是在NYU和FAIR进行博士研究时创建的
PyTorch Lightning是在NYU和FAIR进行博士研究时创建的

PyTorch Lightning是为从事AI研究的专业研究人员博士生而创建的。

Lightning来自我的博士学位。人工智能研究的纽约大学CILVR和Facebook的AI研究。结果,该框架被设计为具有极强的可扩展性,同时又使最先进的AI研究技术(例如TPU训练)变得微不足道。

现在,核心贡献者都在使用Lightning推动AI的发展,并继续添加新的炫酷功能。

但是,简单的界面使专业的生产团队和新手可以使用Pytorch和PyTorch Lightning社区开发的最新技术。

Lightning拥有超过320名贡献者,由11名研究科学家,博士研究生和专业深度学习工程师组成的核心团队。
在这里插入图片描述

2 大纲

本教程将引导您构建一个简单的MNIST分类器,并排显示PyTorch和PyTorch Lightning代码。虽然Lightning可以构建任何任意复杂的系统,使用MNIST来说明如何将PyTorch代码重构为PyTorch Lightning。

完整的代码可在此Colab Notebook中获得。

典型的AI研究项目
在研究项目中,我们通常希望确定以下关键组成部分:

  • 模型 the model(s)
  • 数据 the data
  • 损失 the loss
  • 优化器 the optimizer(s)

2.1 模型的构建

设计一个三层全连接神经网络,该网络以28x28的图像作为输入,并输出10个可能标签上的概率分布。

首先,在PyTorch中定义模型
在这里插入图片描述
该模型定义了计算图,以将MNIST图像作为输入,并将其转换为10个类别(0-9数字)的概率分布。

3-layer network (illustration by: William Falcon)
3-layer network (illustration by: William Falcon)

要将模型转换为PyTorch Lightning,只需将pl.LightningModule替换掉nn.Module
在这里插入图片描述

Lightning 提供了结构化的 PyTorch code

在这里插入图片描述
看!两者的代码完全相同!
这意味着可以像使用PyTorch模块一样完全使用LightningModule,例如预测
在这里插入图片描述
或者用于预训练
在这里插入图片描述

2.2 数据 data

在本教程中,使用MNIST。

在这里插入图片描述
让我们生成MNIST的三个部分,即训练,验证和测试部分。

同样,PyTorch中的代码与Lightning中的代码相同。

数据集被添加到数据加载器Dataloader中,该数据加载器处理数据集的加载shuffling, batching

简而言之,数据准备包括四个步骤:

  1. 下载图片
  2. 图像变换(这些是高度主观的)。
  3. 生成训练,验证和测试数据集拆分。
  4. 将每个数据集拆分包装在DataLoader中

同样,除了将PyTorch代码组织为4个函数之外,代码完全相同:
在这里插入图片描述
对于此代码中的一些关键函数解释如下:

2.3 prepare_data

此功能处理下载和任何数据处理。此功能可确保当您使用多个GPU时,不会下载多个数据集或对数据进行双重操作。

这是因为每个GPU将执行相同的PyTorch,从而导致重复。所有在Lightning的代码可以确保关键部件是从所谓的仅一个GPU。

2.4 dataloader

** train_dataloader,val_dataloadertest_dataloader**

每一个都负责返回适当的数据拆分。Lightning以这种方式进行构造,因此非常清楚如何操作数据。如果曾经阅读用PyTorch编写的随机github代码,则几乎看不到如何操纵数据。

Lightning甚至允许多个数据加载器进行测试或验证。

这段代码是根据我们所谓的DataModule进行组织的。尽管这是100%可选的,并且闪电可以直接使用DataLoaders,但DataModule可以使您的数据可重用并且易于共享。

2.5 优化器 Optimizer

现在选择如何进行优化。将使用Adam而不是SGD,因为它在大多数DL研究中都是很好的默认设置。
在这里插入图片描述
同样,这两者完全相同,只是它被组织到configure optimizers功能中。

Lightning非常容易扩展。例如,如果想使用多个优化器(即GAN),则可以在此处返回两者。
在这里插入图片描述
还会注意到,在Lightning中,传入了self.parameters() 而不是model,因为LightningModule就是model。

2.6 损失函数 loss

对于n向分类,要计算交叉熵损失。交叉熵与将使用的NegativeLogLikelihood(log_softmax)相同。
再次……代码是完全一样的!

2.7 训练和验证循环

Training and Validation Loop

我们汇总了训练所需的所有关键要素:

  • 模型(3层NN) The model (3-layer NN)
  • 数据集(MNIST)
  • 优化器 optimizer
  • 损失 loss

现在,执行一个完整的训练例程,该例程执行以下操作:

  • 迭代多个eooch(一个epoch是对数据集D的完整遍历)
    在这里插入图片描述

  • 每个epoch以批处理b(batch)的小块迭代数据集
    在这里插入图片描述

  • 执行前向通过 forward pass
    在这里插入图片描述

  • 计算损失 loss
    在这里插入图片描述

  • 执行反向传递以计算每个权重的所有梯度
    在这里插入图片描述

  • 优化
    在这里插入图片描述
    在PyTorch和Lightning中,伪代码都看起来像这样
    在这里插入图片描述
    但这就是 Lightning不同的地方。在PyTorch中,自己编写了for循环,这意味着必须记住要以正确的顺序调用正确的东西-这为错误留下了很多空间。

即使模型很简单,也不会像开始做更高级的事情那样,例如使用多个GPU,梯度裁剪,提早停止,检查点,TPU训练,16位精度等……代码复杂性将迅速爆炸。

即使模型很简单,也不会一开始就做更高级的事情

这是PyTorch和Lightning的验证和训练循环
在这里插入图片描述
这就是Lightning的美。它是抽象模板,内容保持不变, 只做了结构上的调整。这意味着您仍在编写PyTorch,除了您的代码结构变得良好。

这增加了可读性,有助于再现性!

2.8 Lightning Trainer

The trainer is how we abstract the boilerplate code.
在这里插入图片描述
同样,这是可能的,因为要做的就是将PyTorch代码组织到LightningModule中

3 完整代码对比

3.1 Full Training Loop for PyTorch

用PyTorch编写的完整MNIST示例如下:

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

# -----------------
# MODEL
# -----------------
class LightningMNISTClassifier(pl.LightningModule):

  def __init__(self):
    super(LightningMNISTClassifier, self).__init__()

    # mnist images are (1, 28, 28) (channels, width, height) 
    self.layer_1 = torch.nn.Linear(28 * 28, 128)
    self.layer_2 = torch.nn.Linear(128, 256)
    self.layer_3 = torch.nn.Linear(256, 10)

  def forward(self, x):
    batch_size, channels, width, height = x.sizes()

    # (b, 1, 28, 28) -> (b, 1*28*28)
    x = x.view(batch_size, -1)

    # layer 1
    x = self.layer_1(x)
    x = torch.relu(x)

    # layer 2
    x = self.layer_2(x)
    x = torch.relu(x)

    # layer 3
    x = self.layer_3(x)

    # probability distribution over labels
    x = torch.log_softmax(x, dim=1)

    return x


# ----------------
# DATA
# ----------------
transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
mnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)
mnist_test = MNIST(os.getcwd(), train=False, download=True, transform=transform)

# train (55,000 images), val split (5,000 images)
mnist_train, mnist_val = random_split(mnist_train, [55000, 5000])
mnist_test = MNIST(os.getcwd(), train=False, download=True)

# The dataloaders handle shuffling, batching, etc...
mnist_train = DataLoader(mnist_train, batch_size=64)
mnist_val = DataLoader(mnist_val, batch_size=64)
mnist_test = DataLoader(mnist_test, batch_size=64)

# ----------------
# OPTIMIZER
# ----------------
pytorch_model = MNISTClassifier()
optimizer = torch.optim.Adam(pytorch_model.parameters(), lr=1e-3)

# ----------------
# LOSS
# ----------------
def cross_entropy_loss(logits, labels):
  return F.nll_loss(logits, labels)

# ----------------
# TRAINING LOOP
# ----------------
num_epochs = 1
for epoch in range(num_epochs):

  # TRAINING LOOP
  for train_batch in mnist_train:
    x, y = train_batch

    logits = pytorch_model(x)
    loss = cross_entropy_loss(logits, y)
    print('train loss: ', loss.item())

    loss.backward()

    optimizer.step()
    optimizer.zero_grad()

  # VALIDATION LOOP
  with torch.no_grad():
    val_loss = []
    for val_batch in mnist_val:
      x, y = val_batch
      logits = pytorch_model(x)
      val_loss.append(cross_entropy_loss(logits, y).item())

    val_loss = torch.mean(torch.tensor(val_loss))
    print('val_loss: ', val_loss.item())

3.2 Full Training loop in Lightning

Lightning中的完整训练循环
Lightning版本完全相同,除了:

  • 核心成分由LightningModule组织
  • 训练者/验证循环代码已由Trainer抽象化
import torch
from torch import nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader, random_split
from torch.nn import functional as F
from torchvision.datasets import MNIST
from torchvision import datasets, transforms
import os

class LightningMNISTClassifier(pl.LightningModule):

  def __init__(self):
    super().__init__()

    # mnist images are (1, 28, 28) (channels, width, height) 
    self.layer_1 = torch.nn.Linear(28 * 28, 128)
    self.layer_2 = torch.nn.Linear(128, 256)
    self.layer_3 = torch.nn.Linear(256, 10)

  def forward(self, x):
      batch_size, channels, width, height = x.size()

      # (b, 1, 28, 28) -> (b, 1*28*28)
      x = x.view(batch_size, -1)

      # layer 1 (b, 1*28*28) -> (b, 128)
      x = self.layer_1(x)
      x = torch.relu(x)

      # layer 2 (b, 128) -> (b, 256)
      x = self.layer_2(x)
      x = torch.relu(x)

      # layer 3 (b, 256) -> (b, 10)
      x = self.layer_3(x)

      # probability distribution over labels
      x = torch.log_softmax(x, dim=1)

      return x

  def cross_entropy_loss(self, logits, labels):
    return F.nll_loss(logits, labels)

  def training_step(self, train_batch, batch_idx):
      x, y = train_batch
      logits = self.forward(x)
      loss = self.cross_entropy_loss(logits, y)
      self.log('train_loss', loss)
      return loss

  def validation_step(self, val_batch, batch_idx):
      x, y = val_batch
      logits = self.forward(x)
      loss = self.cross_entropy_loss(logits, y)
      self.log('val_loss', loss)

  def configure_optimizers(self):
      optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
      return optimizer

# data
# transforms for images
transform=transforms.Compose([transforms.ToTensor(), 
                              transforms.Normalize((0.1307,), (0.3081,))])

# prepare transforms standard to MNIST
mnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)
mnist_test = MNIST(os.getcwd(), train=False, download=True, transform=transform)

train_dataloader = DataLoader(mnist_train, batch_size=64)
val_loader = DataLoader(mnist_test, batch_size=64)

# train
model = LightningMNISTClassifier()
trainer = pl.Trainer()

trainer.fit(model, train_dataloader, val_loader)

3.3 将dataloader写进DataModule中


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

class LightningMNISTClassifier(pl.LightningModule):

  def __init__(self):
    super().__init__()

    # mnist images are (1, 28, 28) (channels, width, height) 
    self.layer_1 = torch.nn.Linear(28 * 28, 128)
    self.layer_2 = torch.nn.Linear(128, 256)
    self.layer_3 = torch.nn.Linear(256, 10)

  def forward(self, x):
      batch_size, channels, width, height = x.size()

      # (b, 1, 28, 28) -> (b, 1*28*28)
      x = x.view(batch_size, -1)

      # layer 1 (b, 1*28*28) -> (b, 128)
      x = self.layer_1(x)
      x = torch.relu(x)

      # layer 2 (b, 128) -> (b, 256)
      x = self.layer_2(x)
      x = torch.relu(x)

      # layer 3 (b, 256) -> (b, 10)
      x = self.layer_3(x)

      # probability distribution over labels
      x = torch.log_softmax(x, dim=1)

      return x

  def cross_entropy_loss(self, logits, labels):
    return F.nll_loss(logits, labels)

  def training_step(self, train_batch, batch_idx):
      x, y = train_batch
      logits = self.forward(x)
      loss = self.cross_entropy_loss(logits, y)
      self.log('train_loss', loss)
      return loss


  def validation_step(self, val_batch, batch_idx):
      x, y = val_batch
      logits = self.forward(x)
      loss = self.cross_entropy_loss(logits, y)
      self.log('val_loss', loss)

  def configure_optimizers(self):
    optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
    return optimizer


class MNISTDataModule(pl.LightningDataModule):

  def setup(self, stage):
    # transforms for images
    transform=transforms.Compose([transforms.ToTensor(), 
                                  transforms.Normalize((0.1307,), (0.3081,))])
      
    # prepare transforms standard to MNIST
    self.mnist_train = MNIST(os.getcwd(), train=True, download=True, transform=transform)
    self.mnist_test = MNIST(os.getcwd(), train=False, download=True, transform=transform)

  def train_dataloader(self):
    return DataLoader(self.mnist_train, batch_size=64)

  def val_dataloader(self):
    return DataLoader(self.mnist_test, batch_size=64)

data_module = MNISTDataModule()

# train
model = LightningMNISTClassifier()
trainer = pl.Trainer()

trainer.fit(model, data_module)

4 Highlights(亮点)

  1. 如果没有Lightning,则可以将PyTorch代码分为任意部分。使用Lightning,这是结构化
  2. 除了在Lightning中进行结构化外,这两者的代码完全相同。(值得大笑🤭)。
  3. 随着项目的复杂性增加,代码将不会因为Lightning提取其中的大部分内容。
  4. 保留了PyTorch的灵活性,因为可以完全控制训练中的关键点。例如,可以使用任意复杂的training_step,例如seq2seq
def training_step(self, batch, batch_idx):
    x, y = batch
    
    # define your own forward and loss calculation
    hidden_states = self.encoder(x)
     
    # even as complex as a seq-2-seq + attn model
    # (this is just a toy, non-working example to illustrate)
    start_token = '<SOS>'
    last_hidden = torch.zeros(...)
    loss = 0
    for step in range(max_seq_len):
        attn_context = self.attention_nn(hidden_states, start_token)
        pred = self.decoder(start_token, attn_context, last_hidden) 
        last_hidden = pred
        pred = self.predict_nn(pred)
        loss += self.loss(last_hidden, y[step])
        
    #toy example as well
    loss = loss / max_seq_len
    return {'loss': loss} 
  1. 在Lightning中,有很多免费赠品,例如进度条
    在这里插入图片描述
    同样还有漂亮的权重总结(weights summary)

在这里插入图片描述
还有tensorboard日志
在这里插入图片描述
还有 checkpointing, and early stopping

5 附加功能

但是Lightning以开箱即用的东西(例如TPU训练等)而闻名。

在Lightning中,可以在CPU,GPU,多个GPU或TPU上训练模型,而无需更改PyTorch代码的一行。

5.1 16位精度训练

Trainer(precision=16)

5.2 多种日志记录方法

使用Tensorboard的其他5种替代方法进行记录点击查看

使用Neptune.AI进行日志记录(鸣谢:Neptune.ai)
在这里插入图片描述
使用Comet.ml记录
在这里插入图片描述
甚至有一个内置的探查器profiler,可以告诉训练中瓶颈的位置。

trainer = Trainer(..., profiler=True)

将此标志设置为True, 将提供如下输出
在这里插入图片描述
或更高级的输出(如果需要)
在这里插入图片描述
在这里插入图片描述

5.3 多GPU 训练

还可以一次在多个GPU上进行训练而无需做任何工作(仍然必须提交SLURM作业)
在这里插入图片描述
它支持大约40种其他功能,可以在文档中阅读这些功能。

5.4 Hooks的可扩展性

可能想知道Lightning如何为做到这一点,又以某种方式做到这一点,以便完全掌控一切?

与keras或其他高级框架不同,Lightning不会隐藏任何必要的细节。但是,如果确实需要自己修改训练的各个方面,那么有两个主要选择。

首先是通过覆盖钩子( hooks)的可扩展性。这是一

以上是关于如何从Pytorch 到 Pytorch Lightning | 简要介绍的主要内容,如果未能解决你的问题,请参考以下文章

从基础概念到实现,小白如何快速入门PyTorch

如何从TensorFlow迁移到PyTorch

PyTorch 完全入门指南!从线性回归逻辑回归到图像分类

PyTorch 1.0 中文官方教程:使用ONNX将模型从PyTorch传输到Caffe2和移动端

PyTorch从入门到精通100讲-PyTorch张量从概念到应用

从 `pytorch-pretrained-bert` 迁移到 `pytorch-transformers` 关于模型()输出的问题