|
|
|
|
公众号矩阵

9个技巧让你的PyTorch模型训练变得飞快!

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

作者:佚名来源:机器学习算法与自然语言处理|2020-11-20 10:40

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

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

我明白,网上都是各种神经网络加速指南,但是一个checklist都没有(现在有了),使用这个清单,一步一步确保你能榨干你模型的所有性能。

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

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

我们会讲到:

  •  使用DataLoaders
  •  DataLoader中的workers数量
  •  Batch size
  •  梯度累计
  •  保留的计算图
  •  移动到单个
  •  16-bit 混合精度训练
  •  移动到多个GPUs中(模型复制)
  •  移动到多个GPU-nodes中 (8+GPUs)
  •  思考模型加速的技巧

Pytorch-Lightning

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

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

  1. from pytorch_lightning import Trainer  
  2. model = LightningModule(…)  
  3. trainer = Trainer()  
  4. trainer.fit(model) 

1. DataLoaders

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

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

  1. dataset = MNIST(root=self.hparams.data_root, traintrain=train, download=True 
  2. loader = DataLoader(dataset, batch_size=32shuffle=True
  3. for batch in loader:  
  4.   x, y = batch  
  5.   model.training_step(x, y)  
  6.   ... 

2. DataLoaders 中的 workers 的数量

另一个加速的神奇之处是允许批量并行加载。因此,您可以一次装载nb_workers个batch,而不是一次装载一个batch。

  1. # slow  
  2. loader = DataLoader(dataset, batch_size=32shuffle=True 
  3. # fast (use 10 workers) 
  4. loader = DataLoader(dataset, batch_size=32shuffle=Truenum_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个前向传播和向后传播,然后再执行一次优化步骤。

  1. # clear last step  
  2. optimizer.zero_grad()  
  3. # 16 accumulated gradient steps  
  4. scaled_loss = 0  
  5. for accumulated_step_i in range(16):  
  6.      out = model.forward()  
  7.      loss = some_loss(out,y)      
  8.      loss.backward()  
  9.       scaled_loss += loss.item()      
  10. # update weights after 8 steps. effective batch = 8*16  
  11. optimizer.step()  
  12. # loss is now scaled up by the number of accumulated batches  
  13. actual_loss = scaled_loss / 16 

在lightning中,全部都给你做好了,只需要设置accumulate_grad_batches=16:

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

5. 保留的计算图

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

  1. losses = []  
  2. ...  
  3. losses.append(loss) 
  4. print(f'current loss: {torch.mean(losses)'}) 

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

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

Lightning会非常小心,确保不会保留计算图的副本。

6. 单个GPU训练

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

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

  1. # put model on GPU  
  2. model.cuda(0)  
  3. # put data on gpu (cuda on a variable returns a cuda copy)  
  4. xx = x.cuda(0)  
  5. # runs on GPU now  
  6. model(x) 

如果你使用Lightning,你什么都不用做,只需要设置Trainer(gpus=1)。

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

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

  1. # expensive  
  2. xx = x.cuda(0)# very expensive  
  3. xx = x.cpu()  
  4. xx = x.cuda(0) 

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

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

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

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

7. 16-bit 精度

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

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

  1. # enable 16-bit on the model and the optimizer  
  2. model, optimizers = amp.initialize(model, optimizers, opt_level='O2' 
  3. # when doing .backward, let amp do it so it can scale the loss  
  4. with amp.scale_loss(loss, optimizer) as scaled_loss:                     
  5.      scaled_loss.backward() 

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

在lightning中,启用16bit并不需要修改模型中的任何内容,也不需要执行我上面所写的操作。设置Trainer(precision=16)就可以了。

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

8. 移动到多个GPUs中

现在,事情变得非常有趣了。有3种(也许更多?)方法来进行多GPU训练。

分batch训练

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

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

  1. # copy model on each GPU and give a fourth of the batch to each  
  2. model = DataParallel(model, devices=[0, 1, 2 ,3])  
  3. # out has 4 outputs (one for each gpu)  
  4. out = model(x.cuda(0)) 

在lightning中,你只需要增加GPUs的数量,然后告诉trainer,其他什么都不用做。

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

模型分布训练

将模型的不同部分放在不同的GPU上,batch按顺序移动

有时你的模型可能太大不能完全放到内存中。例如,带有编码器和解码器的序列到序列模型在生成输出时可能会占用20GB RAM。在本例中,我们希望将编码器和解码器放在独立的GPU上。

  1. # each model is sooo big we can't fit both in memory  
  2. encoder_rnn.cuda(0)  
  3. decoder_rnn.cuda(1)  
  4. # run input through encoder on GPU 0  
  5. encoder_out = encoder_rnn(x.cuda(0))  
  6. # run output through decoder on the next GPU  
  7. out = decoder_rnn(encoder_out.cuda(1))  
  8. # normally we want to bring all outputs back to GPU 0  
  9. outout = out.cuda(0) 

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

  1. class MyModule(LightningModule):  
  2.     def __init__():  
  3.         self.encoder = RNN(...)  
  4.         self.decoder = RNN(...)  
  5.     def forward(x): 
  6.          # models won't be moved after the first forward because   
  7.         # they are already on the correct GPUs  
  8.         self.encoder.cuda(0)  
  9.         self.decoder.cuda(1)  
  10.         out = self.encoder(x)  
  11.         out = self.decoder(out.cuda(1))       
  12.  # don't pass GPUs to trainer 
  13. model = MyModule()  
  14. trainer = Trainer()  
  15. trainer.fit(model) 

两者混合

在上面的情况下,编码器和解码器仍然可以从并行化操作中获益。

  1. # change these lines  
  2. self.encoder = RNN(...)  
  3. self.decoder = RNN(...)  
  4. # to these  
  5. # now each RNN is based on a different gpu set  
  6. self.encoder = DataParallel(self.encoder, devices=[0, 1, 2, 3])  
  7. self.decoder = DataParallel(self.encoder, devices=[4, 5, 6, 7])  
  8. # in forward...  
  9. out = self.encoder(x.cuda(0))  
  10. # notice inputs on first gpu in device  
  11. 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的事情:

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

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

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

Lightning还附带了一个SlurmCluster管理器,可以方便地帮助你提交SLURM作业的正确详细信息。

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

事实证明,distributedDataParallel比DataParallel快得多,因为它只执行梯度同步的通信。所以,一个好的hack是使用distributedDataParallel替换DataParallel,即使是在单机上进行训练。

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

  1. # train on 4 gpus on the same machine MUCH faster than DataParallel  
  2. 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。针对你的具体问题,请查阅相关文献,看看人们都忽略了什么!

【编辑推荐】

  1. 文本或代码中 n 和 r 的区别
  2. 代码不止|Google 高效的前端开发秘籍全披露!
  3. 因未发项目奖金,一名程序员决定删代码泄愤
  4. 一段代码瘫痪一座工厂?奇安信现场展示攻防全过程
  5. 代码不止|一篇看懂 TensorFlow 如何助力高效开发
【责任编辑:庞桂玉 TEL:(010)68476606】

点赞 0
分享:
大家都在看
猜你喜欢

订阅专栏+更多

云原生架构实践

云原生架构实践

新技术引领移动互联网进入急速赛道
共3章 | KaliArch

29人订阅学习

数据中心和VPDN网络建设案例

数据中心和VPDN网络建设案例

漫画+案例
共20章 | 捷哥CCIE

193人订阅学习

搭建数据中心实验Lab

搭建数据中心实验Lab

实验平台Datacenter
共5章 | ITGO(老曾)

119人订阅学习

订阅51CTO邮刊

点击这里查看样刊

订阅51CTO邮刊

51CTO服务号

51CTO官微