加速 PyTorch 模型训练的 9 个技巧(收藏)

2022-04-17 14:08:17
目录
Pytorch-Lightning1.DataLoaders2.DataLoaders中的workers的数量3.Batchsize4.梯度累加5.保留的计算图6.单个GPU训练7.16-bit精度8.移动到多个GPUs中9.多节点GPU训练10.福利!在单个节点上多GPU更快的训练对模型加速的思考

加速PyTorch模型训练的9个技巧(收藏)

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

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

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

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

得batch的一部分。

# copy model on each GPU and give a fourth of the batch to eachmodel = 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 trainingtrainer = Trainer(gpus=[0, 1, 2, 3])trainer.fit(model)

  模型分布训练

加速PyTorch模型训练的9个技巧(收藏)

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

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

# each model is sooo big we can't fit both in memoryencoder_rnn.cuda(0)decoder_rnn.cuda(1) # run input through encoder on GPU 0encoder_out = encoder_rnn(x.cuda(0)) # run output through decoder on the next GPUout = decoder_rnn(encoder_out.cuda(1)) # normally we want to bring all outputs back to GPU 0out = 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 trainermodel = MyModule()trainer = Trainer()trainer.fit(model)

两者混合

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

# change these linesself.encoder = RNN(...)self.decoder = RNN(...) # to these# now each RNN is based on a different gpu setself.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 www.easck.comin devicesout = self.decoder(out.cuda(4))  # <--- the 4 here

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

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

9. 多节点GPU训练

加速PyTorch模型训练的9个技巧(收藏)

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

如果你已经做到了这一步,那么你现在可以在几分钟内训练Imagenet了!这并没有你想象的那么难,但是它可能需要你对计算集群的更多知识。这些说明假设你正在集群上使用SLURM。

Pytorch允许多节点训练,通过在每个节点上复制每个GPU上的模型并同步梯度。所以,每个模型都是在每个GPU上独立初始化的,本质上独立地在数据的一个分区上训练,除了它们都从所有模型接收梯度更新。

在高层次上:

在每个GPU上初始化一个模型的副本(确保设置种子,让每个模型初始化到相同的权重,否则它会失败)。将数据集分割成子集(使用DistributedSampler)。每个GPU只在它自己的小子集上训练。在.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 nodestrainer = 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 DataParalleltrainer = Trainer(distributed_backend='ddp', gpus=[0, 1, 2, 3])

对模型加速的思考

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

我将模型分成几个部分:

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

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

接下来,我试图最大化我的batch size,这通常是受GPU内存大小的限制。现在,需要关注在使用大的batch size的时候如何在多个GPUs上分布并最小化延迟(比如,我可能会尝试着在多个gpu上使用8000 +的有效batch size)。

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