轉(zhuǎn)載于:https://mp.weixin.qq.com/s/WNRz8D9FOlZqTdcjm1usjw
PyTorch Lightning :https://github.com/williamFalcon/pytorch-lightning/projects
用戶文檔:https://williamfalcon.github.io/pytorch-lightning/
1. 介紹
- Lightning是基于Pytorch的一個(gè)光包裝器,它可以幫助研究人員自動(dòng)訓(xùn)練模型,但關(guān)鍵的模型部件還是由研究人員完全控制诡壁。
- 參照此篇教程沟娱,獲得更有力的范例(https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/single_gpu_node_template.py?source=post_page---------------------------)
- Lightning采用最新、最尖端的方法捌锭,將犯錯(cuò)的可能性降到最低。
2. DataLoader
- 這可能是最容易提速的地方∈蹋靠保存h5py或numpy文件來(lái)加速數(shù)據(jù)加載的日子已經(jīng)一去不復(fù)返了。用 Pytorch dataloader加載圖像數(shù)據(jù)非常簡(jiǎn)單教馆。
- 在Lightning中逊谋,你無(wú)需指定一個(gè)訓(xùn)練循環(huán),只需定義dataLoaders土铺,訓(xùn)練器便會(huì)在需要時(shí)調(diào)用它們胶滋。
- https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py?source=post_page---------------------------#L163-L217
3. DataLoaders中的進(jìn)程數(shù)
- 加快速度的第二個(gè)秘訣在于允許批量并行加載。所以悲敷,你可以一次加載許多批量究恤,而不是一次加載一個(gè)。
# slowloader = DataLoader(dataset, batch_size=32, shuffle=True)
2# fast (use 10 workers)loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=10)
4.批尺寸
- 在開(kāi)始下一步優(yōu)化步驟之前后德,將批量大小調(diào)高到CPU內(nèi)存或GPU內(nèi)存允許的最大值部宿。接下來(lái)的部分將著重于減少內(nèi)存占用,這樣就可以繼續(xù)增加批尺寸瓢湃。記住理张,你很可能需要再次更新學(xué)習(xí)率。如果將批尺寸增加一倍绵患,最好將學(xué)習(xí)速度也提高一倍雾叭。
5.累積梯度
- 假如已經(jīng)最大限度地使用了計(jì)算資源,而批尺寸仍然太低(假設(shè)為8)落蝙,那我們則需為梯度下降模擬更大的批尺寸织狐,以供精準(zhǔn)估計(jì)。
- 假設(shè)想讓批尺寸達(dá)到128筏勒。然后移迫,在執(zhí)行單個(gè)優(yōu)化器步驟前,將執(zhí)行16次前向和后向傳播(批量大小為8)管行。
# clear last stepoptimizer.zero_grad()
2
3# 16 accumulated gradient stepsscaled_loss = 0for accumulated_step_i in range(16): out = model.forward() loss = some_loss(out,y) loss.backward()
4
5 scaled_loss += loss.item()
6
7# update weights after 8 steps. effective batch = 8*16optimizer.step()
8
9# loss is now scaled up by the number of accumulated batchesactual_loss = scaled_loss / 16
- 而在Lightning中厨埋,這些已經(jīng)自動(dòng)執(zhí)行了。只需設(shè)置標(biāo)記:
trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)
6. 保留計(jì)算圖
- 撐爆內(nèi)存很簡(jiǎn)單捐顷,只要不釋放指向計(jì)算圖形的指針揽咕,比如……為記錄日志保存loss悲酷。
losses = []
...losses.append(loss)
print(f current loss: {torch.mean(losses) })
- 上述的問(wèn)題在于,loss仍然有一個(gè)圖形副本亲善。在這種情況中设易,可用.item()來(lái)釋放它。
# badlosses.append(loss)
# goodlosses.append(loss.item())
- Lightning會(huì)特別注意蛹头,讓其無(wú)法保留圖形副本
7.轉(zhuǎn)至單GPU
- 一旦完成了前面的步驟顿肺,就可以進(jìn)入GPU訓(xùn)練了。GPU的訓(xùn)練將對(duì)許多GPU核心上的數(shù)學(xué)計(jì)算進(jìn)行并行處理渣蜗。能加速多少取決于使用的GPU類(lèi)型屠尊。個(gè)人使用的話,推薦使用2080Ti耕拷,公司使用的話可用V100讼昆。
- 剛開(kāi)始你可能會(huì)覺(jué)得壓力很大,但其實(shí)只需做兩件事: 1)將你的模型移動(dòng)到GPU上骚烧;2)在用其運(yùn)行數(shù)據(jù)時(shí)浸赫,把數(shù)據(jù)導(dǎo)至GPU中。
# put model on GPUmodel.cuda(0)
# put data on gpu (cuda on a variable returns a cuda copy) x = x.cuda(0)
# runs on GPU nowmodel(x)
- 如果使用Lightning赃绊,則不需要對(duì)代碼做任何操作既峡。只需設(shè)置標(biāo)記
ask lightning to use gpu 0 for trainingtrainer = Trainer(gpus=[0])trainer.fit(model)
- 在GPU進(jìn)行訓(xùn)練時(shí),要注意限制CPU和GPU之間的傳輸量碧查。
# expensivex = x.cuda(0)
# very expensivex = x.cpu()x = x.cuda(0)
- 例如运敢,如果耗盡了內(nèi)存,不要為了省內(nèi)存忠售,將數(shù)據(jù)移回CPU传惠。嘗試用其他方式優(yōu)化代碼,或者在用這種方法之前先跨GPUs分配代碼稻扬。此外還要注意進(jìn)行強(qiáng)制GPUs同步的操作涉枫。例如清除內(nèi)存緩存。
# really bad idea.Stops all the GPUs until they all catch up
torch.cuda.empty_cache()
- 但是如果使用Lightning腐螟,那么只有在定義Lightning模塊時(shí)可能會(huì)出現(xiàn)這種問(wèn)題。Lightning特別注意避免此類(lèi)錯(cuò)誤困后。
8. 16位混合精度訓(xùn)練
- 16位精度可以有效地削減一半的內(nèi)存占用乐纸。大多數(shù)模型都是用32位精度數(shù)進(jìn)行訓(xùn)練的。然而最近的研究發(fā)現(xiàn)摇予,使用16位精度汽绢,模型也可以很好地工作〔啻鳎混合精度指的是宁昭,用16位訓(xùn)練一些特定的模型跌宛,而權(quán)值類(lèi)的用32位訓(xùn)練。
- 要想在Pytorch中用16位精度,先從NVIDIA中安裝 apex 圖書(shū)館并對(duì)你的模型進(jìn)行這些更改积仗。
# enable 16-bit on the model and the optimize rmodel,
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包會(huì)處理大部分事情疆拘。如果梯度爆炸或趨于零,它甚至?xí)U(kuò)大loss寂曹。
- 在Lightning中哎迄, 使用16位很簡(jiǎn)單(https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/?source=post_page---------------------------#16-bit-mixed-precision),不需對(duì)你的模型做任何修改隆圆,也不用完成上述操作漱挚。
trainer = Trainer(amp_level=’O2 , use_amp=False)trainer.fit(model)
9.移至多GPU
-
分批量訓(xùn)練
??第一種方法叫做分批量訓(xùn)練。這一策略將模型復(fù)制到每個(gè)GPU上渺氧,而每個(gè)GPU會(huì)分到該批量的一部分旨涝。
# 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中,可以直接指示訓(xùn)練器增加GPU數(shù)量侣背,而無(wú)需完成上述任何操作白华。
# ask lightning to use 4 GPUs for training
trainer = Trainer(gpus=[0, 1, 2, 3])trainer.fit(model)
-
分模型訓(xùn)練
??將模型的不同部分分配給不同的GPU,按順序分配批量秃踩。有時(shí)模型可能太大衬鱼,內(nèi)存不足以支撐。比如憔杨,帶有編碼器和解碼器的Sequence to Sequence模型在生成輸出時(shí)可能會(huì)占用20gb的內(nèi)存鸟赫。在這種情況下,我們希望把編碼器和解碼器放在單獨(dú)的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
out = encoder_rnn(x.cuda(0))
# run output through decoder on the next GPU
out = decoder_rnn(x.cuda(1))
# normally we want to bring all outputs back to GPU 0
out = out.cuda(0)
- 對(duì)于這種類(lèi)型的訓(xùn)練抛蚤,無(wú)需將Lightning訓(xùn)練器分到任何GPU上。與之相反寻狂,只要把自己的模塊導(dǎo)入正確的GPU的Lightning模塊中:
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)
-
混合兩種訓(xùn)練方法
??在上面的例子中岁经,編碼器和解碼器仍然可以從并行化每個(gè)操作中獲益。我們現(xiàn)在可以更具創(chuàng)造力了蛇券。
# 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 devices
out = self.decoder(out.cuda(4)) # <--- the 4 here
-
使用多GPUs時(shí)需注意的事項(xiàng)
??如果該設(shè)備上已存在model.cuda()缀壤,那么它不會(huì)完成任何操作。
??始終輸入到設(shè)備列表中的第一個(gè)設(shè)備上纠亚。
??跨設(shè)備傳輸數(shù)據(jù)非常昂貴塘慕,不到萬(wàn)不得已不要這樣做。
??優(yōu)化器和梯度將存儲(chǔ)在GPU 0上蒂胞。因此图呢,GPU 0使用的內(nèi)存很可能比其他處理器大得多。
10.轉(zhuǎn)至多GPU階段(8+GPUs)
- 每臺(tái)機(jī)器上的各GPU都可獲取一份模型的副本。每臺(tái)機(jī)器分得一部分?jǐn)?shù)據(jù)蛤织,并僅針對(duì)該部分?jǐn)?shù)據(jù)進(jìn)行訓(xùn)練赴叹。各機(jī)器彼此同步梯度。
- 做到了這一步指蚜,就可以在幾分鐘內(nèi)訓(xùn)練Imagenet數(shù)據(jù)集了! 這沒(méi)有想象中那么難乞巧,但需要更多有關(guān)計(jì)算集群的知識(shí)。這些指令假定你正在集群上使用SLURM姚炕。
- Pytorch在各個(gè)GPU上跨節(jié)點(diǎn)復(fù)制模型并同步梯度摊欠,從而實(shí)現(xiàn)多節(jié)點(diǎn)訓(xùn)練。因此柱宦,每個(gè)模型都是在各GPU上獨(dú)立初始化的些椒,本質(zhì)上是在數(shù)據(jù)的一個(gè)分區(qū)上獨(dú)立訓(xùn)練的,只是它們都接收來(lái)自所有模型的梯度更新掸刊。
-
高級(jí)階段:
??在各GPU上初始化一個(gè)模型的副本(確保設(shè)置好種子免糕,使每個(gè)模型初始化到相同的權(quán)值,否則操作會(huì)失效)忧侧。
??將數(shù)據(jù)集分成子集石窑。每個(gè)GPU只在自己的子集上訓(xùn)練。
??On .backward() 所有副本都會(huì)接收各模型梯度的副本蚓炬。只有此時(shí)松逊,模型之間才會(huì)相互通信。 - Pytorch有一個(gè)很好的抽象概念肯夏,叫做分布式數(shù)據(jù)并行處理经宏,它可以為你完成這一操作。要使用DDP(分布式數(shù)據(jù)并行處理)驯击,需要做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)
- Pytorch團(tuán)隊(duì)對(duì)此有一份詳細(xì)的實(shí)用教程(https://github.com/pytorch/examples/blob/master/imagenet/main.py?source=post_page---------------------------)
- 然而烁兰,在Lightning中,這是一個(gè)自帶功能徊都。只需設(shè)定節(jié)點(diǎn)數(shù)標(biāo)志沪斟,其余的交給Lightning處理就好。
trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7])
- Lightning還附帶了一個(gè)SlurmCluster管理器暇矫,可助你簡(jiǎn)單地提交SLURM任務(wù)的正確細(xì)節(jié)主之。
10.更快的多GPU單節(jié)點(diǎn)訓(xùn)練
- 事實(shí)證明,分布式數(shù)據(jù)并行處理要比數(shù)據(jù)并行快得多,因?yàn)槠湮ㄒ坏耐ㄐ攀翘荻韧健R虼诵古簦詈糜梅植际綌?shù)據(jù)并行處理替換數(shù)據(jù)并行,即使只是在做單機(jī)訓(xùn)練。
- 在Lightning中枉长,通過(guò)將distributed_backend設(shè)置為ddp(分布式數(shù)據(jù)并行處理)并設(shè)置GPU的數(shù)量冀续,這可以很容易實(shí)現(xiàn)琼讽。
# train on 4 gpus on the same machine MUCH faster than DataParallel
trainer = Trainer(distributed_backend= ddp , gpus=[0, 1, 2, 3])