9個(gè)技巧讓你的PyTorch模型訓(xùn)練變得飛快！
不要讓你的神經(jīng)網(wǎng)絡(luò)變成這樣
讓我們面對(duì)現(xiàn)實(shí)吧，你的模型可能還停留在石器時(shí)代。我敢打賭你仍然使用32位精度或GASP甚至只在一個(gè)GPU上訓(xùn)練。
我明白，網(wǎng)上都是各種神經(jīng)網(wǎng)絡(luò)加速指南，但是一個(gè)checklist都沒(méi)有（現(xiàn)在有了），使用這個(gè)清單，一步一步確保你能榨干你模型的所有性能。
本指南從最簡(jiǎn)單的結(jié)構(gòu)到最復(fù)雜的改動(dòng)都有，可以使你的網(wǎng)絡(luò)得到最大的好處。我會(huì)給你展示示例Pytorch代碼以及可以在Pytorch- lightning Trainer中使用的相關(guān)flags，這樣你可以不用自己編寫(xiě)這些代碼！
**這本指南是為誰(shuí)準(zhǔn)備的？**任何使用Pytorch進(jìn)行深度學(xué)習(xí)模型研究的人，如研究人員、博士生、學(xué)者等，我們?cè)谶@里談?wù)摰哪Ｐ涂赡苄枰慊ㄙM(fèi)幾天的訓(xùn)練，甚至是幾周或幾個(gè)月。
我們會(huì)講到：
- 使用DataLoaders
- DataLoader中的workers數(shù)量
- Batch size
- 梯度累計(jì)
- 保留的計(jì)算圖
- 移動(dòng)到單個(gè)
- 16-bit 混合精度訓(xùn)練
- 移動(dòng)到多個(gè)GPUs中（模型復(fù)制）
- 移動(dòng)到多個(gè)GPU-nodes中 (8+GPUs)
- 思考模型加速的技巧
Pytorch-Lightning
你可以在Pytorch的庫(kù)Pytorch- lightning中找到我在這里討論的每一個(gè)優(yōu)化。Lightning是在Pytorch之上的一個(gè)封裝，它可以自動(dòng)訓(xùn)練，同時(shí)讓研究人員完全控制關(guān)鍵的模型組件。Lightning 使用最新的最佳實(shí)踐，并將你可能出錯(cuò)的地方最小化。
我們?yōu)镸NIST定義LightningModel并使用Trainer來(lái)訓(xùn)練模型。
- from pytorch_lightning import Trainer
- model = LightningModule(…)
- trainer = Trainer()
- trainer.fit(model)
1. DataLoaders
這可能是最容易獲得速度增益的地方。保存h5py或numpy文件以加速數(shù)據(jù)加載的時(shí)代已經(jīng)一去不復(fù)返了，使用Pytorch dataloader加載圖像數(shù)據(jù)很簡(jiǎn)單(對(duì)于NLP數(shù)據(jù)，請(qǐng)查看TorchText)。
在lightning中，你不需要指定訓(xùn)練循環(huán)，只需要定義dataLoaders和Trainer就會(huì)在需要的時(shí)候調(diào)用它們。
- dataset = MNIST(root=self.hparams.data_root, traintrain=train, download=True)
- loader = DataLoader(dataset, batch_size=32, shuffle=True)
- for batch in loader:
- x, y = batch
- model.training_step(x, y)
- ...
2. DataLoaders 中的 workers 的數(shù)量
另一個(gè)加速的神奇之處是允許批量并行加載。因此，您可以一次裝載nb_workers個(gè)batch，而不是一次裝載一個(gè)batch。
- # slow
- loader = DataLoader(dataset, batch_size=32, shuffle=True)
- # fast (use 10 workers)
- loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=10)
3. Batch size
在開(kāi)始下一個(gè)優(yōu)化步驟之前，將batch size增大到CPU-RAM或GPU-RAM所允許的最大范圍。
下一節(jié)將重點(diǎn)介紹如何幫助減少內(nèi)存占用，以便你可以繼續(xù)增加batch size。
記住，你可能需要再次更新你的學(xué)習(xí)率。一個(gè)好的經(jīng)驗(yàn)法則是，如果batch size加倍，那么學(xué)習(xí)率就加倍。
4. 梯度累加
在你已經(jīng)達(dá)到計(jì)算資源上限的情況下，你的batch size仍然太小(比如8)，然后我們需要模擬一個(gè)更大的batch size來(lái)進(jìn)行梯度下降，以提供一個(gè)良好的估計(jì)。
假設(shè)我們想要達(dá)到128的batch size大小。我們需要以batch size為8執(zhí)行16個(gè)前向傳播和向后傳播，然后再執(zhí)行一次優(yōu)化步驟。
- # clear last step
- optimizer.zero_grad()
- # 16 accumulated gradient steps
- scaled_loss = 0
- for accumulated_step_i in range(16):
- out = model.forward()
- loss = some_loss(out,y)
- loss.backward()
- scaled_loss += loss.item()
- # update weights after 8 steps. effective batch = 8*16
- optimizer.step()
- # loss is now scaled up by the number of accumulated batches
- actual_loss = scaled_loss / 16
在lightning中，全部都給你做好了，只需要設(shè)置accumulate_grad_batches=16：
- trainer = Trainer(accumulate_grad_batches=16)
- trainer.fit(model)
5. 保留的計(jì)算圖
一個(gè)最簡(jiǎn)單撐爆你的內(nèi)存的方法是為了記錄日志存儲(chǔ)你的loss。
- losses = []
- ...
- losses.append(loss)
- print(f'current loss: {torch.mean(losses)'})
上面的問(wèn)題是，loss仍然包含有整個(gè)圖的副本。在這種情況下，調(diào)用.item()來(lái)釋放它。
- # bad
- losses.append(loss)
- # good
- losses.append(loss.item())
Lightning會(huì)非常小心，確保不會(huì)保留計(jì)算圖的副本。
6. 單個(gè)GPU訓(xùn)練
一旦你已經(jīng)完成了前面的步驟，是時(shí)候進(jìn)入GPU訓(xùn)練了。在GPU上的訓(xùn)練將使多個(gè)GPU cores之間的數(shù)學(xué)計(jì)算并行化。你得到的加速取決于你所使用的GPU類(lèi)型。我推薦個(gè)人用2080Ti，公司用V100。
乍一看，這可能會(huì)讓你不知所措，但你真的只需要做兩件事:1)移動(dòng)你的模型到GPU, 2)每當(dāng)你運(yùn)行數(shù)據(jù)通過(guò)它，把數(shù)據(jù)放到GPU上。
- # put model on GPU
- model.cuda(0)
- # put data on gpu (cuda on a variable returns a cuda copy)
- xx = x.cuda(0)
- # runs on GPU now
- model(x)
如果你使用Lightning，你什么都不用做，只需要設(shè)置Trainer(gpus=1)。
- # ask lightning to use gpu 0 for training
- trainer = Trainer(gpus=[0])
- trainer.fit(model)
在GPU上進(jìn)行訓(xùn)練時(shí)，要注意的主要事情是限制CPU和GPU之間的傳輸次數(shù)。
- # expensive
- xx = x.cuda(0)# very expensive
- xx = x.cpu()
- xx = x.cuda(0)
如果內(nèi)存耗盡，不要將數(shù)據(jù)移回CPU以節(jié)省內(nèi)存。在求助于GPU之前，嘗試以其他方式優(yōu)化你的代碼或GPU之間的內(nèi)存分布。
另一件需要注意的事情是調(diào)用強(qiáng)制GPU同步的操作。清除內(nèi)存緩存就是一個(gè)例子。
- # really bad idea. Stops all the GPUs until they all catch up
- torch.cuda.empty_cache()
但是，如果使用Lightning，惟一可能出現(xiàn)問(wèn)題的地方是在定義Lightning Module時(shí)。Lightning會(huì)特別注意不去犯這類(lèi)錯(cuò)誤。
7. 16-bit 精度
16bit精度是將內(nèi)存占用減半的驚人技術(shù)。大多數(shù)模型使用32bit精度數(shù)字進(jìn)行訓(xùn)練。然而，最近的研究發(fā)現(xiàn)，16bit模型也可以工作得很好?；旌暇纫馕吨鴮?duì)某些內(nèi)容使用16bit，但將權(quán)重等內(nèi)容保持在32bit。
要在Pytorch中使用16bit精度，請(qǐng)安裝NVIDIA的apex庫(kù)，并對(duì)你的模型進(jìn)行這些更改。
- # 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包會(huì)處理好大部分事情。如果梯度爆炸或趨向于0，它甚至?xí)s放loss。
在lightning中，啟用16bit并不需要修改模型中的任何內(nèi)容，也不需要執(zhí)行我上面所寫(xiě)的操作。設(shè)置Trainer(precision=16)就可以了。
- trainer = Trainer(amp_level='O2', use_amp=False)
- trainer.fit(model)
8. 移動(dòng)到多個(gè)GPUs中
現(xiàn)在，事情變得非常有趣了。有3種(也許更多？)方法來(lái)進(jìn)行多GPU訓(xùn)練。
分batch訓(xùn)練
A) 拷貝模型到每個(gè)GPU中，B) 給每個(gè)GPU一部分batch
第一種方法被稱(chēng)為“分batch訓(xùn)練”。該策略將模型復(fù)制到每個(gè)GPU上，每個(gè)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的數(shù)量，然后告訴trainer，其他什么都不用做。
- # ask lightning to use 4 GPUs for training
- trainer = Trainer(gpus=[0, 1, 2, 3])
- trainer.fit(model)
模型分布訓(xùn)練
將模型的不同部分放在不同的GPU上，batch按順序移動(dòng)
有時(shí)你的模型可能太大不能完全放到內(nèi)存中。例如，帶有編碼器和解碼器的序列到序列模型在生成輸出時(shí)可能會(huì)占用20GB RAM。在本例中，我們希望將編碼器和解碼器放在獨(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
- 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
- outout = out.cuda(0)
對(duì)于這種類(lèi)型的訓(xùn)練，在Lightning中不需要指定任何GPU，你應(yīng)該把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
使用多個(gè)GPU時(shí)要考慮的注意事項(xiàng)：
- 如果模型已經(jīng)在GPU上了，model.cuda()不會(huì)做任何事情。
- 總是把輸入放在設(shè)備列表中的第一個(gè)設(shè)備上。
- 在設(shè)備之間傳輸數(shù)據(jù)是昂貴的，把它作為最后的手段。
- 優(yōu)化器和梯度會(huì)被保存在GPU 0上，因此，GPU 0上使用的內(nèi)存可能會(huì)比其他GPU大得多
9. 多節(jié)點(diǎn)GPU訓(xùn)練
每臺(tái)機(jī)器上的每個(gè)GPU都有一個(gè)模型的副本。每臺(tái)機(jī)器獲得數(shù)據(jù)的一部分，并且只在那部分上訓(xùn)練。每臺(tái)機(jī)器都能同步梯度。
如果你已經(jīng)做到了這一步，那么你現(xiàn)在可以在幾分鐘內(nèi)訓(xùn)練Imagenet了！這并沒(méi)有你想象的那么難，但是它可能需要你對(duì)計(jì)算集群的更多知識(shí)。這些說(shuō)明假設(shè)你正在集群上使用SLURM。
Pytorch允許多節(jié)點(diǎn)訓(xùn)練，通過(guò)在每個(gè)節(jié)點(diǎn)上復(fù)制每個(gè)GPU上的模型并同步梯度。所以，每個(gè)模型都是在每個(gè)GPU上獨(dú)立初始化的，本質(zhì)上獨(dú)立地在數(shù)據(jù)的一個(gè)分區(qū)上訓(xùn)練，除了它們都從所有模型接收梯度更新。
在高層次上：
- 在每個(gè)GPU上初始化一個(gè)模型的副本(確保設(shè)置種子，讓每個(gè)模型初始化到相同的權(quán)重，否則它會(huì)失敗)。
- 將數(shù)據(jù)集分割成子集(使用DistributedSampler)。每個(gè)GPU只在它自己的小子集上訓(xùn)練。
- 在.backward()上，所有副本都接收到所有模型的梯度副本。這是模型之間唯一一次的通信。
Pytorch有一個(gè)很好的抽象，叫做DistributedDataParallel，它可以幫你實(shí)現(xiàn)這個(gè)功能。要使用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, worldworld_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中，只需設(shè)置節(jié)點(diǎn)數(shù)量，它就會(huì)為你處理其余的事情。
- # train on 1024 gpus across 128 nodes
- trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7])
Lightning還附帶了一個(gè)SlurmCluster管理器，可以方便地幫助你提交SLURM作業(yè)的正確詳細(xì)信息。
10. 福利！在單個(gè)節(jié)點(diǎn)上多GPU更快的訓(xùn)練
事實(shí)證明，distributedDataParallel比DataParallel快得多，因?yàn)樗粓?zhí)行梯度同步的通信。所以，一個(gè)好的hack是使用distributedDataParallel替換DataParallel，即使是在單機(jī)上進(jìn)行訓(xùn)練。
在Lightning中，這很容易通過(guò)將distributed_backend設(shè)置為ddp和設(shè)置GPUs的數(shù)量來(lái)實(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])
對(duì)模型加速的思考
盡管本指南將為你提供了一系列提高網(wǎng)絡(luò)速度的技巧，但我還是要給你解釋一下如何通過(guò)查找瓶頸來(lái)思考問(wèn)題。
我將模型分成幾個(gè)部分：
首先，我要確保在數(shù)據(jù)加載中沒(méi)有瓶頸。為此，我使用了我所描述的現(xiàn)有數(shù)據(jù)加載解決方案，但是如果沒(méi)有一種解決方案滿(mǎn)足你的需要，請(qǐng)考慮離線(xiàn)處理和緩存到高性能數(shù)據(jù)存儲(chǔ)中，比如h5py。
接下來(lái)看看你在訓(xùn)練步驟中要做什么。確保你的前向傳播速度快，避免過(guò)多的計(jì)算以及最小化CPU和GPU之間的數(shù)據(jù)傳輸。最后，避免做一些會(huì)降低GPU速度的事情(本指南中有介紹)。
接下來(lái)，我試圖最大化我的batch size，這通常是受GPU內(nèi)存大小的限制?，F(xiàn)在，需要關(guān)注在使用大的batch size的時(shí)候如何在多個(gè)GPUs上分布并最小化延遲（比如，我可能會(huì)嘗試著在多個(gè)gpu上使用8000 +的有效batch size）。
然而，你需要小心大的batch size。針對(duì)你的具體問(wèn)題，請(qǐng)查閱相關(guān)文獻(xiàn)，看看人們都忽略了什么！
