你應(yīng)該知道的18個(gè)PyTorch小技巧

導(dǎo)讀

你應(yīng)該知道的18個(gè)PyTorch小技巧。

調(diào)試深度學(xué)習(xí)的pipelines就像找到最合適的齒輪組合

你為什么要讀這篇文章？

深度學(xué)習(xí)模型的訓(xùn)練/推理過(guò)程涉及很多步驟。在有限的時(shí)間和資源條件下，每個(gè)迭代的速度越快，整個(gè)模型的預(yù)測(cè)性能就越快。我收集了幾個(gè)PyTorch技巧，以最大化內(nèi)存使用效率和最小化運(yùn)行時(shí)間。為了更好地利用這些技巧，我們還需要理解它們?nèi)绾我约盀槭裁从行А?/span>

我首先提供一個(gè)完整的列表和一些代碼片段，這樣你就可以開(kāi)始優(yōu)化你的腳本了。然后我一個(gè)一個(gè)地詳細(xì)地研究它們。對(duì)于每個(gè)技巧，我還提供了代碼片段和注釋?zhuān)嬖V你它是特定于設(shè)備類(lèi)型(CPU/GPU)還是模型類(lèi)型。

列表:

• 數(shù)據(jù)加載

  1、把數(shù)據(jù)放到SSD中

  2、Dataloader(dataset, num_workers=4*num_GPU)

  3、Dataloader(dataset, pin_memory=True)

• 數(shù)據(jù)操作

  4、直接在設(shè)備中創(chuàng)建torch.Tensor，不要在一個(gè)設(shè)備中創(chuàng)建再移動(dòng)到另一個(gè)設(shè)備中

  5、避免CPU和GPU之間不必要的數(shù)據(jù)傳輸

  6、使用torch.from_numpy(numpy_array)或者torch.as_tensor(others)

  7、在數(shù)據(jù)傳輸操作可以重疊時(shí)，使用tensor.to(non_blocking=True)

  8、使用PyTorch JIT將元素操作融合到單個(gè)kernel中。

• 模型結(jié)構(gòu)

  9、在使用混合精度的FP16時(shí)，對(duì)于所有不同架構(gòu)設(shè)計(jì)，設(shè)置尺寸為8的倍數(shù)

• 訓(xùn)練

  10、將batch size設(shè)置為8的倍數(shù)，最大化GPU內(nèi)存的使用

  11、前向的時(shí)候使用混合精度（后向的使用不用）

  12、在優(yōu)化器更新權(quán)重之前，設(shè)置梯度為None，model.zero_grad(set_to_none=True)

  13、梯度積累：每隔x個(gè)batch更新一次權(quán)重，模擬大batch size的效果

• 推理/驗(yàn)證

  14、關(guān)閉梯度計(jì)算

• CNN (卷積神經(jīng)網(wǎng)絡(luò)) 特有的

  15、torch.backends.cudnn.benchmark = True

  16、對(duì)于4D NCHW Tensors，使用channels_last的內(nèi)存格式

  17、在batch normalization之前的卷積層可以去掉bias

• 分布式

  18、用DistributedDataParallel代替DataParallel

第7、11、12、13的代碼片段

#CombiningthetipsNo.7,11,12,13:nonblocking,AMP,setting
#gradientsasNone,andlargereffectivebatchsize
model.train()

#ResetthegradientstoNone
optimizer.zero_grad(set_to_none=True)
scaler=GradScaler()
fori,(features,target)inenumerate(dataloader):
#thesetwocallsarenonblockingandoverlapping
features=features.to('cuda:0',non_blocking=True)
target=target.to('cuda:0',non_blocking=True)

#Forwardpasswithmixedprecision
withtorch.cuda.amp.autocast():#autocastasacontextmanager
output=model(features)
loss=criterion(output,target)

#Backwardpasswithoutmixedprecision
#It'snotrecommendedtousemixedprecisionforbackwardpass
#Becauseweneedmorepreciseloss
scaler.scale(loss).backward()

#Onlyupdateweightseveryother2iterations
#Effectivebatchsizeisdoubled
if(i+1)%2==0or(i+1)==len(dataloader):
#scaler.step()firstunscalesthegradients.
#IfthesegradientscontaininfsorNaNs,
#optimizer.step()isskipped.
scaler.step(optimizer)

#Ifoptimizer.step()wasskipped,
#scalingfactorisreducedbythebackoff_factor
#inGradScaler()
scaler.update()

#ResetthegradientstoNone
optimizer.zero_grad(set_to_none=True)

指導(dǎo)思想

總的來(lái)說(shuō)，你可以通過(guò)3個(gè)關(guān)鍵點(diǎn)來(lái)優(yōu)化時(shí)間和內(nèi)存使用。首先，盡可能減少i/o(輸入/輸出)，使模型管道更多的用于計(jì)算，而不是用于i/o(帶寬限制或內(nèi)存限制)。這樣，我們就可以利用GPU及其他專(zhuān)用硬件來(lái)加速這些計(jì)算。第二，盡量重疊過(guò)程，以節(jié)省時(shí)間。第三，最大限度地提高內(nèi)存使用效率，節(jié)約內(nèi)存。然后，節(jié)省內(nèi)存可以啟用更大的batch size大小，從而節(jié)省更多的時(shí)間。擁有更多的時(shí)間有助于更快的模型開(kāi)發(fā)周期，并導(dǎo)致更好的模型性能。

1、把數(shù)據(jù)移動(dòng)到SSD中

有些機(jī)器有不同的硬盤(pán)驅(qū)動(dòng)器，如HHD和SSD。建議將項(xiàng)目中使用的數(shù)據(jù)移動(dòng)到SSD(或具有更好i/o的硬盤(pán)驅(qū)動(dòng)器)以獲得更快的速度。

2. 在加載數(shù)據(jù)和數(shù)據(jù)增強(qiáng)的時(shí)候異步處理

num_workers=0使數(shù)據(jù)加載需要在訓(xùn)練完成后或前一個(gè)處理已完成后進(jìn)行。設(shè)置num_workers>0有望加快速度，特別是對(duì)于大數(shù)據(jù)的i/o和增強(qiáng)。具體到GPU，有實(shí)驗(yàn)發(fā)現(xiàn)num_workers = 4*num_GPU 具有最好的性能。也就是說(shuō)，你也可以為你的機(jī)器測(cè)試最佳的num_workers。需要注意的是，高num_workers將會(huì)有很大的內(nèi)存消耗開(kāi)銷(xiāo)，這也是意料之中的，因?yàn)楦嗟臄?shù)據(jù)副本正在內(nèi)存中同時(shí)處理。

Dataloader(dataset,num_workers=4*num_GPU)

3. 使用pinned memory來(lái)降低數(shù)據(jù)傳輸

設(shè)置pin_memory=True可以跳過(guò)從可分頁(yè)memory到pinned memory的數(shù)據(jù)傳輸

GPU無(wú)法直接從CPU的可分頁(yè)內(nèi)存中訪問(wèn)數(shù)據(jù)。設(shè)置pin_memory=True 可以為CPU主機(jī)上的數(shù)據(jù)直接分配臨時(shí)內(nèi)存，節(jié)省將數(shù)據(jù)從可分頁(yè)內(nèi)存轉(zhuǎn)移到臨時(shí)內(nèi)存(即固定內(nèi)存又稱(chēng)頁(yè)面鎖定內(nèi)存)的時(shí)間。該設(shè)置可以與num_workers = 4*num_GPU結(jié)合使用。

Dataloader(dataset,pin_memory=True)

4. 直接在設(shè)備中創(chuàng)建張量

只要你需要torch.Tensor，首先嘗試在要使用它們的設(shè)備上創(chuàng)建它們。不要使用原生Python或NumPy創(chuàng)建數(shù)據(jù)，然后將其轉(zhuǎn)換為torch.Tensor。在大多數(shù)情況下，如果你要在GPU中使用它們，直接在GPU中創(chuàng)建它們。

#Randomnumbersbetween0and1
#Sameasnp.random.rand([10,5])
tensor=torch.rand([10,5],device=torch.device('cuda:0'))

#Randomnumbersfromnormaldistributionwithmean0andvariance1
#Sameasnp.random.randn([10,5])
tensor=torch.randn([10,5],device=torch.device('cuda:0'))

唯一的語(yǔ)法差異是NumPy中的隨機(jī)數(shù)生成需要額外的random，例如：np.random.rand() vs torch.rand()。許多其他函數(shù)在NumPy中也有相應(yīng)的函數(shù)：

torch.empty(),torch.zeros(),torch.full(),torch.ones(),torch.eye(),torch.randint(),torch.rand(),torch.randn()

5. 避免在CPU和GPU中傳輸數(shù)據(jù)

正如我在指導(dǎo)思想中提到的，我們希望盡可能地減少I(mǎi)/O。注意下面這些命令：

#BAD!AVOIDTHEMIFUNNECESSARY!
print(cuda_tensor)
cuda_tensor.cpu()
cuda_tensor.to_device('cpu')
cpu_tensor.cuda()
cpu_tensor.to_device('cuda')
cuda_tensor.item()
cuda_tensor.numpy()
cuda_tensor.nonzero()
cuda_tensor.tolist()

#PythoncontrolflowwhichdependsonoperationresultsofCUDAtensors
if(cuda_tensor!=0).all():
run_func()

6. 使用 torch.from_numpy(numpy_array)和torch.as_tensor(others)代替 torch.tensor

torch.tensor() 會(huì)拷貝數(shù)據(jù)

如果源設(shè)備和目標(biāo)設(shè)備都是CPU，torch.from_numpy和torch.as_tensor不會(huì)創(chuàng)建數(shù)據(jù)拷貝。如果源數(shù)據(jù)是NumPy數(shù)組，使用torch.from_numpy(numpy_array) 會(huì)更快。如果源數(shù)據(jù)是一個(gè)具有相同數(shù)據(jù)類(lèi)型和設(shè)備類(lèi)型的張量，那么torch.as_tensor(others) 可以避免拷貝數(shù)據(jù)。others 可以是Python的list， tuple，或者torch.tensor。如果源設(shè)備和目標(biāo)設(shè)備不同，那么我們可以使用下一個(gè)技巧。

torch.from_numpy(numpy_array)
torch.as_tensor(others)

7. 在數(shù)據(jù)傳輸有重疊時(shí)使用tensor.to(non_blocking=True)

重疊數(shù)據(jù)傳輸以減少運(yùn)行時(shí)間

本質(zhì)上，non_blocking=True允許異步數(shù)據(jù)傳輸以減少執(zhí)行時(shí)間。

forfeatures,targetinloader:
#thesetwocallsarenonblockingandoverlapping
features=features.to('cuda:0',non_blocking=True)
target=target.to('cuda:0',non_blocking=True)

#Thisisasynchronizationpoint
#Itwillwaitforprevioustwolines
output=model(features)

8. 使用PyTorch JIT將點(diǎn)操作融合到單個(gè)kernel中

點(diǎn)操作包括常見(jiàn)的數(shù)學(xué)操作，通常是內(nèi)存受限的。PyTorch JIT會(huì)自動(dòng)將相鄰的點(diǎn)操作融合到一個(gè)內(nèi)核中，以保存多次內(nèi)存讀/寫(xiě)操作。例如，通過(guò)將5個(gè)核融合成1個(gè)核，gelu函數(shù)可以被加速4倍。

@torch.jit.script#JITdecorator
deffused_gelu(x):
returnx*0.5*(1.0+torch.erf(x/1.41421))

9 & 10. 在使用混合精度的FP16時(shí)，對(duì)于所有不同架構(gòu)設(shè)計(jì)，設(shè)置圖像尺寸和batch size為8的倍數(shù)

為了最大限度地提高GPU的計(jì)算效率，最好保證不同的架構(gòu)設(shè)計(jì)(包括神經(jīng)網(wǎng)絡(luò)的輸入輸出尺寸/維數(shù)/通道數(shù)和batch size大小)是8的倍數(shù)甚至更大的2的冪(如64、128和最大256)。這是因?yàn)楫?dāng)矩陣的維數(shù)與2的冪倍數(shù)對(duì)齊時(shí)，Nvidia gpu的張量核心(Tensor Cores)在矩陣乘法方面可以獲得最佳性能。矩陣乘法是最常用的操作，也可能是瓶頸，所以它是我們能確保張量/矩陣/向量的維數(shù)能被2的冪整除的最好方法(例如，8、64、128，最多256)。

這些實(shí)驗(yàn)顯示設(shè)置輸出維度和batch size大小為8的倍數(shù)，比如(33712、4088、4096)相比33708，batch size為4084或者4095這些不能被8整除的數(shù)可以加速計(jì)算1.3倍到 4倍。加速度大小取決于過(guò)程類(lèi)型(例如，向前傳遞或梯度計(jì)算)和cuBLAS版本。特別是，如果你使用NLP，請(qǐng)記住檢查輸出維度，這通常是詞匯表大小。

使用大于256的倍數(shù)不會(huì)增加更多的好處，但也沒(méi)有害處。這些設(shè)置取決于cuBLAS和cuDNN版本以及GPU架構(gòu)。你可以在文檔中找到矩陣維數(shù)的特定張量核心要求。由于目前PyTorch AMP多使用FP16，而FP16需要8的倍數(shù)，所以通常推薦使用8的倍數(shù)。如果你有更高級(jí)的GPU，比如A100，那么你可以選擇64的倍數(shù)。如果你使用的是AMD GPU，你可能需要檢查AMD的文檔。

除了將batch size大小設(shè)置為8的倍數(shù)外，我們還將batch size大小最大化，直到它達(dá)到GPU的內(nèi)存限制。這樣，我們可以用更少的時(shí)間來(lái)完成一個(gè)epoch。

11. 在前向中使用混合精度后向中不使用

有些操作不需要float64或float32的精度。因此，將操作設(shè)置為較低的精度可以節(jié)省內(nèi)存和執(zhí)行時(shí)間。對(duì)于各種應(yīng)用，英偉達(dá)報(bào)告稱(chēng)具有Tensor Cores的GPU的混合精度可以提高3.5到25倍的速度。

值得注意的是，通常矩陣越大，混合精度加速度越高。在較大的神經(jīng)網(wǎng)絡(luò)中(例如BERT)，實(shí)驗(yàn)表明混合精度可以加快2.75倍的訓(xùn)練，并減少37%的內(nèi)存使用。具有Volta, Turing, Ampere或Hopper架構(gòu)的較新的GPU設(shè)備(例如，T4, V100, RTX 2060, 2070, 2080, 2080 Ti, A100, RTX 3090, RTX 3080，和RTX 3070)可以從混合精度中受益更多，因?yàn)樗麄冇蠺ensor Core架構(gòu)，它相比CUDA cores有特殊的優(yōu)化。

帶有Tensor Core的NVIDIA架構(gòu)支持不同的精度

值得一提的是，采用Hopper架構(gòu)的H100預(yù)計(jì)將于2022年第三季度發(fā)布，支持FP8 (float8)。PyTorch AMP可能會(huì)支持FP8(目前v1.11.0還不支持FP8)。

在實(shí)踐中，你需要在模型精度性能和速度性能之間找到一個(gè)最佳點(diǎn)。我之前確實(shí)發(fā)現(xiàn)混合精度可能會(huì)降低模型的精度，這取決于算法，數(shù)據(jù)和問(wèn)題。

使用自動(dòng)混合精度(AMP)很容易在PyTorch中利用混合精度。PyTorch中的默認(rèn)浮點(diǎn)類(lèi)型是float32。AMP將通過(guò)使用float16來(lái)進(jìn)行一組操作(例如，matmul, linear, conv2d)來(lái)節(jié)省內(nèi)存和時(shí)間。AMP會(huì)自動(dòng)cast到float32的一些操作(例如，mse_loss, softmax等)。有些操作(例如add)可以操作最寬的輸入類(lèi)型。例如，如果一個(gè)變量是float32，另一個(gè)變量是float16，那么加法結(jié)果將是float32。

autocast自動(dòng)應(yīng)用精度到不同的操作。因?yàn)閾p失和梯度是按照f(shuō)loat16精度計(jì)算的，當(dāng)它們太小時(shí)，梯度可能會(huì)“下溢”并變成零。GradScaler通過(guò)將損失乘以一個(gè)比例因子來(lái)防止下溢，根據(jù)比例損失計(jì)算梯度，然后在優(yōu)化器更新權(quán)重之前取消梯度的比例。如果縮放因子太大或太小，并導(dǎo)致inf或NaN，則縮放因子將在下一個(gè)迭代中更新縮放因子。

scaler=GradScaler()
forfeatures,targetindata:
#Forwardpasswithmixedprecision
withtorch.cuda.amp.autocast():#autocastasacontextmanager
output=model(features)
loss=criterion(output,target)

#Backwardpasswithoutmixedprecision
#It'snotrecommendedtousemixedprecisionforbackwardpass
#Becauseweneedmorepreciseloss
scaler.scale(loss).backward()

#scaler.step()firstunscalesthegradients.
#IfthesegradientscontaininfsorNaNs,
#optimizer.step()isskipped.
scaler.step(optimizer)

#Ifoptimizer.step()wasskipped,
#scalingfactorisreducedbythebackoff_factorinGradScaler()
scaler.update()

你也可以使用autocast 作為前向傳遞函數(shù)的裝飾器。

classAutocastModel(nn.Module):
...
@autocast()#autocastasadecorator
defforward(self,input):
x=self.model(input)
returnx

12. 在優(yōu)化器更新權(quán)重之前將梯度設(shè)置為None

通過(guò)model.zero_grad()或optimizer.zero_grad()將對(duì)所有參數(shù)執(zhí)行memset ，并通過(guò)讀寫(xiě)操作更新梯度。但是，將梯度設(shè)置為None將不會(huì)執(zhí)行memset，并且將使用“只寫(xiě)”操作更新梯度。因此，設(shè)置梯度為None更快。

#Resetgradientsbeforeeachstepofoptimizer
forparaminmodel.parameters():
param.grad=None

#or(PyTorch>=1.7)
model.zero_grad(set_to_none=True)

#or(PyTorch>=1.7)
optimizer.zero_grad(set_to_none=True)

13. 梯度累積：每隔x個(gè)batch再更新梯度，模擬大batch size

這個(gè)技巧是關(guān)于從更多的數(shù)據(jù)樣本積累梯度，以便對(duì)梯度的估計(jì)更準(zhǔn)確，權(quán)重更新更接近局部/全局最小值。這在batch size較小的情況下更有幫助(由于GPU內(nèi)存限制較小或每個(gè)樣本的數(shù)據(jù)量較大)。

fori,(features,target)inenumerate(dataloader):
#Forwardpass
output=model(features)
loss=criterion(output,target)

#Backwardpass
loss.backward()

#Onlyupdateweightseveryother2iterations
#Effectivebatchsizeisdoubled
if(i+1)%2==0or(i+1)==len(dataloader):
#Updateweights
optimizer.step()
#ResetthegradientstoNone
optimizer.zero_grad(set_to_none=True)

14. 在推理和驗(yàn)證的時(shí)候禁用梯度計(jì)算

實(shí)際上，如果只計(jì)算模型的輸出，那么梯度計(jì)算對(duì)于推斷和驗(yàn)證步驟并不是必需的。PyTorch使用一個(gè)中間內(nèi)存緩沖區(qū)來(lái)處理requires_grad=True變量中涉及的操作。因此，如果我們知道不需要任何涉及梯度的操作，通過(guò)禁用梯度計(jì)算來(lái)進(jìn)行推斷/驗(yàn)證，就可以避免使用額外的資源。

#torch.no_grad()asacontextmanager:
withtorch.no_grad():
output=model(input)

#torch.no_grad()asafunctiondecorator:
@torch.no_grad()
defvalidation(model,input):
output=model(input)
returnoutput

15. torch.backends.cudnn.benchmark = True

在訓(xùn)練循環(huán)之前設(shè)置torch.backends.cudnn.benchmark = True可以加速計(jì)算。由于計(jì)算不同內(nèi)核大小卷積的cuDNN算法的性能不同，自動(dòng)調(diào)優(yōu)器可以運(yùn)行一個(gè)基準(zhǔn)來(lái)找到最佳算法。當(dāng)你的輸入大小不經(jīng)常改變時(shí)，建議開(kāi)啟這個(gè)設(shè)置。如果輸入大小經(jīng)常改變，那么自動(dòng)調(diào)優(yōu)器就需要太頻繁地進(jìn)行基準(zhǔn)測(cè)試，這可能會(huì)損害性能。它可以將向前和向后傳播速度提高1.27x到1.70x。

torch.backends.cudnn.benchmark=True

16. 對(duì)于4D NCHW Tensors使用通道在最后的內(nèi)存格式

4D NCHW重新組織成 NHWC格式

使用channels_last內(nèi)存格式以逐像素的方式保存圖像，作為內(nèi)存中最密集的格式。原始4D NCHW張量在內(nèi)存中按每個(gè)通道(紅/綠/藍(lán))順序存儲(chǔ)。轉(zhuǎn)換之后，x = x.to(memory_format=torch.channels_last)，數(shù)據(jù)在內(nèi)存中被重組為NHWC (channels_last格式)。你可以看到RGB層的每個(gè)像素更近了。據(jù)報(bào)道，這種NHWC格式與FP16的AMP一起使用可以獲得8%到35%的加速。

目前，它仍處于beta測(cè)試階段，僅支持4D NCHW張量和一組模型(例如，alexnet，mnasnet家族，mobilenet_v2，resnet家族，shufflenet_v2，squeezenet1，vgg家族)。但我可以肯定，這將成為一個(gè)標(biāo)準(zhǔn)的優(yōu)化。

N,C,H,W=10,3,32,32

x=torch.rand(N,C,H,W)

#Strideisthegapbetweenoneelementtothenextone
#inadimension.
print(x.stride())

#(3072,1024,32,1)#ConvertthetensortoNHWCinmemory
x2=x.to(memory_format=torch.channels_last)

print(x2.shape)#(10,3,32,32)asdimensionsorderpreserved
print(x2.stride())#(3072,1,96,3),whicharesmaller
print((x==x2).all())#Truebecausethevalueswerenotchanged

17. 在batch normalization之前禁用卷積層的bias

這是可行的，因?yàn)樵跀?shù)學(xué)上，bias可以通過(guò)batch normalization的均值減法來(lái)抵消。我們可以節(jié)省模型參數(shù)、運(yùn)行時(shí)的內(nèi)存。

nn.Conv2d(...,bias=False)

18. 使用 DistributedDataParallel代替DataParallel

對(duì)于多GPU來(lái)說(shuō)，即使只有單個(gè)節(jié)點(diǎn)，也總是優(yōu)先使用 DistributedDataParallel而不是 DataParallel ，因?yàn)?DistributedDataParallel 應(yīng)用于多進(jìn)程，并為每個(gè)GPU創(chuàng)建一個(gè)進(jìn)程，從而繞過(guò)Python全局解釋器鎖(GIL)并提高速度。

總結(jié)

在這篇文章中，我列出了一個(gè)清單，并提供了18個(gè)PyTorch技巧的代碼片段。然后，我逐一解釋了它們?cè)诓煌矫娴?a href="http://wenjunhu.com/v/tag/773/" target="_blank">工作原理和原因，包括數(shù)據(jù)加載、數(shù)據(jù)操作、模型架構(gòu)、訓(xùn)練、推斷、cnn特定的優(yōu)化和分布式計(jì)算。一旦你深入理解了它們的工作原理，你可能會(huì)找到適用于任何深度學(xué)習(xí)框架中的深度學(xué)習(xí)建模的通用原則。

審核編輯 ：李倩