PyTorch模型训练实战技巧，突破速度瓶颈

点击关注上方 “ AI深度学习视线 ”，并“ 星标 ”公号

做者：William Falcon

编译：ronghuaiyang 来源：AI公园

导读

一个step by step的指南，很是的实用。

不要让你的神经网络变成这样

让咱们面对现实吧，你的模型可能还停留在石器时代。我敢打赌你仍然使用32位精度或GASP甚至只在一个GPU上训练。

我明白，网上都是各类神经网络加速指南，可是一个checklist都没有（如今有了），使用这个清单，一步一步确保你能榨干你模型的全部性能。

本指南从最简单的结构到最复杂的改动都有，可使你的网络获得最大的好处。我会给你展现示例Pytorch代码以及能够在Pytorch- lightning Trainer中使用的相关flags，这样你能够不用本身编写这些代码！

**这本指南是为谁准备的？**任何使用Pytorch进行深度学习模型研究的人，如研究人员、博士生、学者等，咱们在这里谈论的模型可能须要你花费几天的训练，甚至是几周或几个月。

咱们会讲到：

1. 使用DataLoaders
2. DataLoader中的workers数量
3. Batch size
4. 梯度累计
5. 保留的计算图
6. 移动到单个
7. 16-bit 混合精度训练
8. 移动到多个GPUs中（模型复制）
9. 移动到多个GPU-nodes中 (8+GPUs)
10. 思考模型加速的技巧

Pytorch-Lightning

你能够在Pytorch的库Pytorch- lightning中找到我在这里讨论的每个优化。Lightning是在Pytorch之上的一个封装，它能够自动训练，同时让研究人员彻底控制关键的模型组件。Lightning 使用最新的最佳实践，并将你可能出错的地方最小化。

咱们为MNIST定义LightningModel并使用Trainer来训练模型。

from pytorch_lightning import Trainer
model = LightningModule(…)
trainer = Trainer()
trainer.fit(model)

1. DataLoaders

这多是最容易得到速度增益的地方。保存h5py或numpy文件以加速数据加载的时代已经一去不复返了，使用Pytorch dataloader加载图像数据很简单(对于NLP数据，请查看TorchText)。

在lightning中，你不须要指定训练循环，只须要定义dataLoaders和Trainer就会在须要的时候调用它们。

dataset = MNIST(root=self.hparams.data_root, train=train, download=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
  x, y = batch
  model.training_step(x, y)
  ...

2. DataLoaders 中的 workers 的数量

另外一个加速的神奇之处是容许批量并行加载。所以，您能够一次装载nb_workers个batch，而不是一次装载一个batch。

# slow
loader = DataLoader(dataset, batch_size=32, shuffle=True)
# fast (use 10 workers)
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=10)

3. Batch size

在开始下一个优化步骤以前，将batch size增大到CPU-RAM或GPU-RAM所容许的最大范围。

下一节将重点介绍如何帮助减小内存占用，以便你能够继续增长batch size。

记住，你可能须要再次更新你的学习率。一个好的经验法则是，若是batch size加倍，那么学习率就加倍。

4. 梯度累加

在你已经达到计算资源上限的状况下，你的batch size仍然过小(好比8)，而后咱们须要模拟一个更大的batch size来进行梯度降低，以提供一个良好的估计。

假设咱们想要达到128的batch size大小。咱们须要以batch size为8执行16个前向传播和向后传播，而后再执行一次优化步骤。

# clear last step
optimizer.zero_grad()

# 16 accumulated gradient steps
scaled_loss = 0
for accumulated_step_i in range(16):
     out = model.forward()
     loss = some_loss(out,y)    
     loss.backward()
      scaled_loss += loss.item()
      
# update weights after 8 steps. effective batch = 8*16
optimizer.step()

# loss is now scaled up by the number of accumulated batches
actual_loss = scaled_loss / 16

在lightning中，所有都给你作好了，只须要设置accumulate_grad_batches=16：

trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)

5. 保留的计算图

一个最简单撑爆你的内存的方法是为了记录日志存储你的loss。

losses = []
...
losses.append(loss)

print(f'current loss: {torch.mean(losses)'})

上面的问题是，loss仍然包含有整个图的副本。在这种状况下，调用.item()来释放它。

![1_CER3v8cok2UOBNsmnBrzPQ](9 Tips For Training Lightning-Fast Neural Networks In Pytorch.assets/1_CER3v8cok2UOBNsmnBrzPQ.gif)# bad
losses.append(loss)

# good
losses.append(loss.item())

Lightning会很是当心，确保不会保留计算图的副本。

6. 单个GPU训练

一旦你已经完成了前面的步骤，是时候进入GPU训练了。在GPU上的训练将使多个GPU cores之间的数学计算并行化。你获得的加速取决于你所使用的GPU类型。我推荐我的用2080Ti，公司用V100。

乍一看，这可能会让你不知所措，但你真的只须要作两件事:1)移动你的模型到GPU, 2)每当你运行数据经过它，把数据放到GPU上。

# put model on GPU
model.cuda(0)

# put data on gpu (cuda on a variable returns a cuda copy)
x = x.cuda(0)

# runs on GPU now
model(x)

若是你使用Lightning，你什么都不用作，只须要设置Trainer(gpus=1)。

# ask lightning to use gpu 0 for training
trainer = Trainer(gpus=[0])
trainer.fit(model)

在GPU上进行训练时，要注意的主要事情是限制CPU和GPU之间的传输次数。

# expensive
x = x.cuda(0)# very expensive
x = x.cpu()
x = x.cuda(0)

若是内存耗尽，不要将数据移回CPU以节省内存。在求助于GPU以前，尝试以其余方式优化你的代码或GPU之间的内存分布。

另外一件须要注意的事情是调用强制GPU同步的操做。清除内存缓存就是一个例子。

# really bad idea. Stops all the GPUs until they all catch up
torch.cuda.empty_cache()

可是，若是使用Lightning，唯一可能出现问题的地方是在定义Lightning Module时。Lightning会特别注意不去犯这类错误。

7. 16-bit 精度

16bit精度是将内存占用减半的惊人技术。大多数模型使用32bit精度数字进行训练。然而，最近的研究发现，16bit模型也能够工做得很好。混合精度意味着对某些内容使用16bit，但将权重等内容保持在32bit。

要在Pytorch中使用16bit精度，请安装NVIDIA的apex库，并对你的模型进行这些更改。

# enable 16-bit on the model and the optimizer
model, optimizers = amp.initialize(model, optimizers, opt_level='O2')

# when doing .backward, let amp do it so it can scale the loss
with amp.scale_loss(loss, optimizer) as scaled_loss:                      
    scaled_loss.backward()

amp包会处理好大部分事情。若是梯度爆炸或趋向于0，它甚至会缩放loss。

在lightning中，启用16bit并不须要修改模型中的任何内容，也不须要执行我上面所写的操做。设置Trainer(precision=16)就能够了。

trainer = Trainer(amp_level='O2', use_amp=False)
trainer.fit(model)

8. 移动到多个GPUs中

如今，事情变得很是有趣了。有3种(也许更多？)方法来进行多GPU训练。

分batch训练

A) 拷贝模型到每一个GPU中，B) 给每一个GPU一部分batch

第一种方法被称为“分batch训练”。该策略将模型复制到每一个GPU上，每一个GPU得到batch的一部分。

# copy model on each GPU and give a fourth of the batch to each
model = DataParallel(model, devices=[0, 1, 2 ,3])

# out has 4 outputs (one for each gpu)
out = model(x.cuda(0))

在lightning中，你只须要增长GPUs的数量，而后告诉trainer，其余什么都不用作。

# ask lightning to use 4 GPUs for training
trainer = Trainer(gpus=[0, 1, 2, 3])
trainer.fit(model)

模型分布训练

将模型的不一样部分放在不一样的GPU上，batch按顺序移动

有时你的模型可能太大不能彻底放到内存中。例如，带有编码器和解码器的序列到序列模型在生成输出时可能会占用20GB RAM。在本例中，咱们但愿将编码器和解码器放在独立的GPU上。

# each model is sooo big we can't fit both in memory
encoder_rnn.cuda(0)
decoder_rnn.cuda(1)

# run input through encoder on GPU 0
encoder_out = encoder_rnn(x.cuda(0))

# run output through decoder on the next GPU
out = decoder_rnn(encoder_out.cuda(1))

# normally we want to bring all outputs back to GPU 0
out = out.cuda(0)

对于这种类型的训练，在Lightning中不须要指定任何GPU，你应该把LightningModule中的模块放到正确的GPU上。

class MyModule(LightningModule):
    def __init__():
        self.encoder = RNN(...)
        self.decoder = RNN(...)
    def forward(x):
        # models won't be moved after the first forward because 
        # they are already on the correct GPUs
        self.encoder.cuda(0)
        self.decoder.cuda(1)
        out = self.encoder(x)
        out = self.decoder(out.cuda(1))
        
# don't pass GPUs to trainer
model = MyModule()
trainer = Trainer()
trainer.fit(model)

二者混合

在上面的状况下，编码器和解码器仍然能够从并行化操做中获益。

# change these lines
self.encoder = RNN(...)
self.decoder = RNN(...)

# to these
# now each RNN is based on a different gpu set
self.encoder = DataParallel(self.encoder, devices=[0, 1, 2, 3])
self.decoder = DataParallel(self.encoder, devices=[4, 5, 6, 7])

# in forward...
out = self.encoder(x.cuda(0))

# notice inputs on first gpu in device
sout = self.decoder(out.cuda(4))  # <--- the 4 here

使用多个GPU时要考虑的注意事项：

• 若是模型已经在GPU上了，model.cuda()不会作任何事情。
• 老是把输入放在设备列表中的第一个设备上。
• 在设备之间传输数据是昂贵的，把它做为最后的手段。
• 优化器和梯度会被保存在GPU 0上，所以，GPU 0上使用的内存可能会比其余GPU大得多。

9. 多节点GPU训练

每台机器上的每一个GPU都有一个模型的副本。每台机器得到数据的一部分，而且只在那部分上训练。每台机器都能同步梯度。

若是你已经作到了这一步，那么你如今能够在几分钟内训练Imagenet了！这并无你想象的那么难，可是它可能须要你对计算集群的更多知识。这些说明假设你正在集群上使用SLURM。

Pytorch容许多节点训练，经过在每一个节点上复制每一个GPU上的模型并同步梯度。因此，每一个模型都是在每一个GPU上独立初始化的，本质上独立地在数据的一个分区上训练，除了它们都从全部模型接收梯度更新。

在高层次上：

1. 在每一个GPU上初始化一个模型的副本(确保设置种子，让每一个模型初始化到相同的权重，不然它会失败)。
2. 将数据集分割成子集(使用DistributedSampler)。每一个GPU只在它本身的小子集上训练。
3. 在.backward()上，全部副本都接收到全部模型的梯度副本。这是模型之间惟一一次的通讯。

Pytorch有一个很好的抽象，叫作DistributedDataParallel，它能够帮你实现这个功能。要使用DDP，你须要作4的事情：

def tng_dataloader():
     d = MNIST()
     
     # 4: Add distributed sampler
     # sampler sends a portion of tng data to each machine
     dist_sampler = DistributedSampler(dataset)
     dataloader = DataLoader(d, shuffle=False, sampler=dist_sampler)
     
def main_process_entrypoint(gpu_nb):
     # 2: set up connections  between all gpus across all machines
     # all gpus connect to a single GPU "root"
     # the default uses env://
     world = nb_gpus * nb_nodes
     dist.init_process_group("nccl", rank=gpu_nb, world_size=world)
     
     # 3: wrap model in DPP
     torch.cuda.set_device(gpu_nb)
     model.cuda(gpu_nb)
     model = DistributedDataParallel(model, device_ids=[gpu_nb])
     
     # train your model now...
     
if  __name__ == '__main__':
     # 1: spawn number of processes
     # your cluster will call main for each machine
     mp.spawn(main_process_entrypoint, nprocs=8)

然而，在Lightning中，只需设置节点数量，它就会为你处理其他的事情。

# train on 1024 gpus across 128 nodes
trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7])

Lightning还附带了一个SlurmCluster管理器，能够方便地帮助你提交SLURM做业的正确详细信息。

10. 福利！在单个节点上多GPU更快的训练

事实证实，distributedDataParallel比DataParallel快得多，由于它只执行梯度同步的通讯。因此，一个好的hack是使用distributedDataParallel替换DataParallel，即便是在单机上进行训练。

在Lightning中，这很容易经过将distributed_backend设置为ddp和设置GPUs的数量来实现。

# train on 4 gpus on the same machine MUCH faster than DataParallel
trainer = Trainer(distributed_backend='ddp', gpus=[0, 1, 2, 3])

对模型加速的思考

尽管本指南将为你提供了一系列提升网络速度的技巧，但我仍是要给你解释一下如何经过查找瓶颈来思考问题。

我将模型分红几个部分：

首先，我要确保在数据加载中没有瓶颈。为此，我使用了我所描述的现有数据加载解决方案，可是若是没有一种解决方案知足你的须要，请考虑离线处理和缓存到高性能数据存储中，好比h5py。

接下来看看你在训练步骤中要作什么。确保你的前向传播速度快，避免过多的计算以及最小化CPU和GPU之间的数据传输。最后，避免作一些会下降GPU速度的事情(本指南中有介绍)。

接下来，我试图最大化个人batch size，这一般是受GPU内存大小的限制。如今，须要关注在使用大的batch size的时候如何在多个GPUs上分布并最小化延迟（好比，我可能会尝试着在多个gpu上使用8000 +的有效batch size）。

然而，你须要当心大的batch size。针对你的具体问题，请查阅相关文献，看看人们都忽略了什么！

—END—

英文原文：https://towardsdatascience.com/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565

推荐阅读

目标检测：

• 干货 | YOLOV5 训练自动驾驶数据集，并转Tensorrt，收藏！

• 超越YOLOv4-tiny！YOLObile：移动设备上的实时目标检测

• 52.1 AP！RepPoints V2：用于目标检测的验证+回归，刚刚开源！
  

• 50.4mAP(已开源) | FCOSv2.0强势归来！Anchor-free家族再添猛将(长文解析)

• COCO 迎来新榜首！DetectoRS以54.7mAP成就目前最高精度检测网络

• 52.1 %AP！AutoAssign：Anchor-free最强检测网络！密集目标检测的可微标签分配

• mAP提高40%！YOLO3改进版—— Poly-YOLO：更快，更精确的检测和实例分割

• YOLOv4 详解版！一句话总结：速度差很少的精度碾压，精度差很少的速度碾压！

• ... 更多精彩内容请进入公号查看
  

目标分割：

• 75.5%mIoU / 58FPS！FANet：带有快速注意力的实时语义分割
  

• CVPR 2020 | CenterMask : Anchor-Free 实时实例分割(长文详解)

• ECCV2020 | 300+FPS！超快的结构感知车道检测新网络

• 50.1%mAP | D2Det：高质量目标检测和实例分割新高度！(文末送书！)

• 图像分割 | CCNet: Criss-Cross Attention for Semantic Segmentation

• 【长文】图像分割 |  YOLACT：开创性的one-stage实时实例分割检测器，名字致敬YOLO

• 投稿CVPR 2021 | BiSeNet V2出来了！72.6%的mIOU, 156FPS的速度！让分割飞起来！

• ... 更多精彩内容请进入公号查看

潮水自会来去，但心志得坚若磐石。即使成不了那根定海神针，也至少不是那随意被拍上岸的野鬼游魂。 

长按关注 "AI深度视线"
最新人工智能、深度学习、SLAM干货奉上！

本文分享自微信公众号 - AI深度学习视线（AI_DeepSight）。
若有侵权，请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一块儿分享。