9个技巧让你的PyTorch模型训练变得飞快!
时间:2022-08-03 17:19:00 sitemap
点击上方“小白学视觉”,选择加"星标"或“置顶”
重磅干货,第一时间送达 本文转自|视觉算法
不要让你的神经网络变成这样
让我们面对现实。你的模型可能还停留在石器时代。我敢打赌你仍然使用32位精度或GASP甚至只有一个GPU上训练。
我知道互联网是各种神经网络加速指南,但一个checklist没有(现在有了),用这个清单,一步一步保证你能榨干型的所有性能。
从最简单的结构到最复杂的变化,这个指南可以让你的网络受益最大。我会给你一个例子Pytorch代码和可在Pytorch- lightning Trainer使用中的相关性flags,这样你就不用自己写这些代码了!
这个指南是为谁准备的?任何使用Pytorch对于研究人员、博士生、学者等进行深度学习模型研究的人,我们在这里讨论的模型可能需要你几天的训练,甚至几周或几个月。
我们会说:
使用DataLoaders
DataLoader中的workers数量
Batch size
梯度累计
保留的计算图
移动到单个
16-bit 混合精度训练
移动到多个GPUs中(模型复制)
移动到多个GPU-nodes中 (8 GPUs)
思考模型加速的技巧
Pytorch-Lighting
你可以在Pytorch的库Pytorch- lightning在这里找到我讨论的每一个优化。Lightning是在Pytorch上述包装可自动训练,使研究人员完全控制关键模型组件。Lightning 使用最新的最佳实践,最小化你可能犯的错误。
我们为MNIST定义LightningModel并使用Trainer训练模型。
frompytorch_lightningimportTrainer 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) forbatchinloader: 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(use10workers) 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为了提供良好的估计,梯度下降。
假设我们想达到128batch size大小。我们需要以batch size对8进行16个前向传播和后向传播,然后进行优化步骤。
#clearlaststep optimizer.zero_grad() #16accumulatedgradientsteps scaled_loss=0 foraccumulated_step_iinrange(16): out=model.forward() loss=some_loss(out,y) loss.backward() scaled_loss =loss.item() #updateweightsafter8steps.effectivebatch=8*16 optimizer.step() #lossisnowscaledupbythenumberofaccumulatedbatches 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'currentloss:{torch.mean(losses)'})
上述问题是,loss它仍然包含整个图片的副本。在这种情况下,调用.item()释放它。
#bad losses.append(loss) #good losses.append(loss.item())
Lightning确保计算图的副本不会保留。
一旦你已经完成了前面的步骤,是时候进入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上独立初始化的,本质上独立地在数据的一个分区上训练,除了它们都从所有模型接收梯度更新。
在高层次上:
在每个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 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。针对你的具体问题,请查阅相关文献,看看人们都忽略了什么!
英文原文:https://towardsdatascience.com/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565
好消息,小白学视觉团队的知识星球开通啦,为了感谢大家的支持与厚爱,团队决定将价值149元的知识星球现时免费加入。各位小伙伴们要抓住机会哦!
下载1:OpenCV-Contrib扩展模块中文版教程
在「小白学视觉」公众号后台回复:扩展模块中文教程,即可下载全网第一份OpenCV扩展模块教程中文版,涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。
下载2:Python视觉实战项目52讲
在「小白学视觉」公众号后台回复:Python视觉实战项目,即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目,助力快速学校计算机视觉。
下载3:OpenCV实战项目20讲
在「小白学视觉」公众号后台回复:OpenCV实战项目20讲,即可下载含有20个基于OpenCV实现20个实战项目,实现OpenCV学习进阶。
交流群
欢迎加入公众号读者群一起和同行交流,目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群(以后会逐渐细分),请扫描下面微信号加群,备注:”昵称+学校/公司+研究方向“,例如:”张三 + 上海交大 + 视觉SLAM“。请按照格式备注,否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告,否则会请出群,谢谢理解~