如何让PyTorch模型训练变得飞快?

让我们面对现实吧,你的模型可能还停留在石器时代。我敢打赌你仍然使用32位精度或gasp甚至只在一个gpu上训练。
我明白,网上都是各种神经网络加速指南,但是一个checklist都没有(现在有了),使用这个清单,一步一步确保你能榨干你模型的所有性能。
本指南从最简单的结构到最复杂的改动都有,可以使你的网络得到最大的好处。我会给你展示示例pytorch代码以及可以在pytorch- lightning trainer中使用的相关flags,这样你可以不用自己编写这些代码!
这篇指南是为谁准备的?任何使用pytorch进行深度学习模型研究的人,如研究人员、博士生、学者等,我们在这里谈论的模型可能需要你花费几天的训练,甚至是几周或几个月。
我们会讲到:
使用dataloaders
dataloader中的workers数量
batch size
梯度累计
保留的计算图
移动到单个
16-bit 混合精度训练
移动到多个gpus中(模型复制)
移动到多个gpu-nodes中 (8+gpus)
思考模型加速的技巧
你可以在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来进行梯度下降,以提供一个良好的估计。
假设我们想要达到128的batch 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()来释放它。
![1_cer3v8cok2uobnsmnbrzpq](9tipsfortraininglightning-fastneuralnetworksinpytorch.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上。
#putmodelongpu model.cuda(0) #putdataongpu(cudaonavariablereturnsacudacopy) x=x.cuda(0) #runsongpunow model(x)
如果你使用lightning,你什么都不用做,只需要设置trainer(gpus=1)。
#asklightningtousegpu0fortraining trainer=trainer(gpus=[0]) trainer.fit(model)
在gpu上进行训练时,要注意的主要事情是限制cpu和gpu之间的传输次数。
#expensive x=x.cuda(0)#veryexpensive x=x.cpu() x=x.cuda(0)
如果内存耗尽,不要将数据移回cpu以节省内存。在求助于gpu之前,尝试以其他方式优化你的代码或gpu之间的内存分布。
另一件需要注意的事情是调用强制gpu同步的操作。清除内存缓存就是一个例子。
#reallybadidea.stopsallthegpusuntiltheyallcatchup torch.cuda.empty_cache()
但是,如果使用lightning,惟一可能出现问题的地方是在定义lightning module时。lightning会特别注意不去犯这类错误。
7. 16-bit 精度
16bit精度是将内存占用减半的惊人技术。大多数模型使用32bit精度数字进行训练。然而,最近的研究发现,16bit模型也可以工作得很好。混合精度意味着对某些内容使用16bit,但将权重等内容保持在32bit。
要在pytorch中使用16bit精度,请安装nvidia的apex库,并对你的模型进行这些更改。
#enable16-bitonthemodelandtheoptimizer model,optimizers=amp.initialize(model,optimizers,opt_level='o2') #whendoing.backward,letampdoitsoitcanscaletheloss withamp.scale_loss(loss,optimizer)asscaled_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的一部分。
#copymodeloneachgpuandgiveafourthofthebatchtoeach model=dataparallel(model,devices=[0,1,2,3]) #outhas4outputs(oneforeachgpu) out=model(x.cuda(0))
在lightning中,你只需要增加gpus的数量,然后告诉trainer,其他什么都不用做。
#asklightningtouse4gpusfortraining trainer=trainer(gpus=[0,1,2,3]) trainer.fit(model)
模型分布训练
将模型的不同部分放在不同的gpu上,batch按顺序移动
有时你的模型可能太大不能完全放到内存中。例如,带有编码器和解码器的序列到序列模型在生成输出时可能会占用20gb ram。在本例中,我们希望将编码器和解码器放在独立的gpu上。
#eachmodelissooobigwecan'tfitbothinmemory encoder_rnn.cuda(0) decoder_rnn.cuda(1) #runinputthroughencoderongpu0 encoder_out=encoder_rnn(x.cuda(0)) #runoutputthroughdecoderonthenextgpu out=decoder_rnn(encoder_out.cuda(1)) #normallywewanttobringalloutputsbacktogpu0 out=out.cuda(0)
对于这种类型的训练,在lightning中不需要指定任何gpu,你应该把lightningmodule中的模块放到正确的gpu上。
classmymodule(lightningmodule): def__init__(): self.encoder=rnn(...) self.decoder=rnn(...) defforward(x): #modelswon'tbemovedafterthefirstforwardbecause #theyarealreadyonthecorrectgpus self.encoder.cuda(0) self.decoder.cuda(1) out=self.encoder(x) out=self.decoder(out.cuda(1)) #don'tpassgpustotrainer model=mymodule() trainer=trainer() trainer.fit(model)
两者混合
在上面的情况下,编码器和解码器仍然可以从并行化操作中获益。
#changetheselines self.encoder=rnn(...) self.decoder=rnn(...) #tothese #noweachrnnisbasedonadifferentgpuset self.encoder=dataparallel(self.encoder,devices=[0,1,2,3]) self.decoder=dataparallel(self.encoder,devices=[4,5,6,7]) #inforward... out=self.encoder(x.cuda(0)) #noticeinputsonfirstgpuindevice sout=self.decoder(out.cuda(4))#<---the4here
使用多个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的事情:
deftng_dataloader(): d=mnist() #4:adddistributedsampler #samplersendsaportionoftngdatatoeachmachine dist_sampler=distributedsampler(dataset) dataloader=dataloader(d,shuffle=false,sampler=dist_sampler) defmain_process_entrypoint(gpu_nb): #2:setupconnectionsbetweenallgpusacrossallmachines #allgpusconnecttoasinglegpuroot #thedefaultusesenv:// world=nb_gpus*nb_nodes dist.init_process_group(nccl,rank=gpu_nb,world_size=world) #3:wrapmodelindpp torch.cuda.set_device(gpu_nb) model.cuda(gpu_nb) model=distributeddataparallel(model,device_ids=[gpu_nb]) #trainyourmodelnow... if__name__=='__main__': #1:spawnnumberofprocesses #yourclusterwillcallmainforeachmachine mp.spawn(main_process_entrypoint,nprocs=8)
然而,在lightning中,只需设置节点数量,它就会为你处理其余的事情。
#trainon1024gpusacross128nodes 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的数量来实现。
#trainon4gpusonthesamemachinemuchfasterthandataparallel 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。针对你的具体问题,请查阅相关文献,看看人们都忽略了什么!
原文标题:9个技巧让你的pytorch模型训练变得飞快!
文章出处:【微信公众号:深度学习自然语言处理】欢迎添加关注!文章转载请注明出处。


ai文件用什么软件打开_ai文件怎么打开
AI大屏旗舰荣耀Note10:6.95英寸AMOLED全面屏配以双独立音腔
为什么物联网非常重要
限流器是啥 常见的限流器
可穿戴设备现在是噱头,未来是主流
如何让PyTorch模型训练变得飞快?
接错三相相序机电设备可能瞬间报废?
天马加大国产OLED产量 或扼杀三星霸主地位
双摄对决:华为p9和荣耀8评测,怎么选?
安波福在可持续发展领域获颁多项荣誉
英特尔宋继强抗疫思考:科学为源,技术为器,促进更加有意义的科技发现与探索
华大电子推出物联网通信专用M2M安全芯片
RS触发器的应用场景
vivo iQOO苏宁手机品鉴会,全民焕新节重磅新品亮相苏宁全球总部
差压式流量计流量显示小的原因及解决办法
新唐科技W588C015主板介绍
汽车“缺芯”使欧洲半导体业捉襟见肘
价值创造是工业互联网从概念走向实践的关键
中国移动的5G服务还远吗
三星S23系列前9个月销量2506万部 较S22系列同期增长23%