優化 PyTorch 速度和內存效率的技巧彙總

作者：ronghuaiyang

調試深度學習的 pipelines 就像找到最合適的齒輪組合

你爲什麼要讀這篇文章？

深度學習模型的訓練 / 推理過程涉及很多步驟。在有限的時間和資源條件下，每個迭代的速度越快，整個模型的預測性能就越快。我收集了幾個 PyTorch 技巧，以最大化內存使用效率和最小化運行時間。爲了更好地利用這些技巧，我們還需要理解它們如何以及爲什麼有效。

我首先提供一個完整的列表和一些代碼片段，這樣你就可以開始優化你的腳本了。然後我一個一個地詳細地研究它們。對於每個技巧，我還提供了代碼片段和註釋，告訴你它是特定於設備類型 (CPU/GPU) 還是模型類型。

列表:

• 數據加載

  1、把數據放到 SSD 中

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

  3、Dataloader(dataset, pin_memory=True)

• 數據操作

  4、直接在設備中創建torch.Tensor，不要在一個設備中創建再移動到另一個設備中

  5、避免 CPU 和 GPU 之間不必要的數據傳輸

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

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

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

• 模型結構

  9、在使用混合精度的 FP16 時，對於所有不同架構設計，設置尺寸爲 8 的倍數

• 訓練

  10、將 batch size 設置爲 8 的倍數，最大化 GPU 內存的使用

  11、前向的時候使用混合精度（後向的使用不用）

  12、在優化器更新權重之前，設置梯度爲None，model.zero_grad(set_to_none=True)

  13、梯度積累：每隔 x 個 batch 更新一次權重，模擬大 batch size 的效果

• 推理 / 驗證

  14、關閉梯度計算

• CNN (卷積神經網絡) 特有的

  15、torch.backends.cudnn.benchmark = True

  16、對於 4D NCHW Tensors，使用 channels_last 的內存格式

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

• 分佈式

  18、用DistributedDataParallel代替DataParallel

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

# Combining the tips No.7, 11, 12, 13: nonblocking, AMP, setting 
# gradients as None, and larger effective batch size
model.train()

# Reset the gradients to None
optimizer.zero_grad(set_to_none=True)
scaler = GradScaler()
for i, (features, target) in enumerate(dataloader):
    # these two calls are nonblocking and overlapping
    features = features.to('cuda:0', non_blocking=True)
    target = target.to('cuda:0', non_blocking=True)
    
    # Forward pass with mixed precision
    with torch.cuda.amp.autocast(): # autocast as a context manager
        output = model(features)
        loss = criterion(output, target)
    
    # Backward pass without mixed precision
    # It's not recommended to use mixed precision for backward pass
    # Because we need more precise loss
    scaler.scale(loss).backward()
    
    # Only update weights every other 2 iterations
    # Effective batch size is doubled
    if (i+1) % 2 == 0 or (i+1) == len(dataloader):
        # scaler.step() first unscales the gradients .
        # If these gradients contain infs or NaNs, 
        # optimizer.step() is skipped.
        scaler.step(optimizer)
        
        # If optimizer.step() was skipped,
        # scaling factor is reduced by the backoff_factor 
        # in GradScaler()
        scaler.update()
        
        # Reset the gradients to None
        optimizer.zero_grad(set_to_none=True)

指導思想

總的來說，你可以通過 3 個關鍵點來優化時間和內存使用。首先，儘可能減少 i/o(輸入 / 輸出)，使模型管道更多的用於計算，而不是用於 i/o(帶寬限制或內存限制)。這樣，我們就可以利用 GPU 及其他專用硬件來加速這些計算。第二，儘量重疊過程，以節省時間。第三，最大限度地提高內存使用效率，節約內存。然後，節省內存可以啓用更大的 batch size 大小，從而節省更多的時間。擁有更多的時間有助於更快的模型開發週期，並導致更好的模型性能。

1、把數據移動到 SSD 中

有些機器有不同的硬盤驅動器，如 HHD 和 SSD。建議將項目中使用的數據移動到 SSD(或具有更好 i/o 的硬盤驅動器) 以獲得更快的速度。

2.  在加載數據和數據增強的時候異步處理

num_workers=0使數據加載需要在訓練完成後或前一個處理已完成後進行。設置num_workers>0 有望加快速度，特別是對於大數據的 i/o 和增強。具體到 GPU，有實驗發現num_workers = 4*num_GPU 具有最好的性能。也就是說，你也可以爲你的機器測試最佳的num_workers。需要注意的是，高num_workers將會有很大的內存消耗開銷，這也是意料之中的，因爲更多的數據副本正在內存中同時處理。

Dataloader(dataset, num_workers=4*num_GPU)

3. 使用 pinned memory 來降低數據傳輸

設置 pin_memory=True 可以跳過從可分頁 memory 到 pinned memory 的數據傳輸

GPU 無法直接從 CPU 的可分頁內存中訪問數據。設置pin_memory=True 可以爲 CPU 主機上的數據直接分配臨時內存，節省將數據從可分頁內存轉移到臨時內存 (即固定內存又稱頁面鎖定內存) 的時間。該設置可以與num_workers = 4*num_GPU結合使用。

Dataloader(dataset, pin_memory=True)

4. 直接在設備中創建張量

只要你需要torch.Tensor，首先嚐試在要使用它們的設備上創建它們。不要使用原生 Python 或 NumPy 創建數據，然後將其轉換爲torch.Tensor。在大多數情況下，如果你要在 GPU 中使用它們，直接在 GPU 中創建它們。

# Random numbers between 0 and 1
# Same as np.random.rand([10,5])
tensor = torch.rand([10, 5], device=torch.device('cuda:0'))

# Random numbers from normal distribution with mean 0 and variance 1
# Same as np.random.randn([10,5])
tensor = torch.randn([10, 5], device=torch.device('cuda:0'))

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

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

5. 避免在 CPU 和 GPU 中傳輸數據

正如我在指導思想中提到的，我們希望儘可能地減少 I/O。注意下面這些命令：

# BAD! AVOID THEM IF UNNECESSARY!
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()

# Python control flow which depends on operation results of CUDA tensors
if (cuda_tensor != 0).all():
    run_func()

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

torch.tensor() 會拷貝數據

如果源設備和目標設備都是 CPU，torch.from_numpy和torch.as_tensor不會創建數據拷貝。如果源數據是 NumPy 數組，使用torch.from_numpy(numpy_array) 會更快。如果源數據是一個具有相同數據類型和設備類型的張量，那麼torch.as_tensor(others) 可以避免拷貝數據。others 可以是 Python 的list， tuple，或者torch.tensor。如果源設備和目標設備不同，那麼我們可以使用下一個技巧。

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

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

重疊數據傳輸以減少運行時間

本質上，non_blocking=True允許異步數據傳輸以減少執行時間。

for features, target in loader:
    # these two calls are nonblocking and overlapping
    features = features.to('cuda:0', non_blocking=True)
    target = target.to('cuda:0', non_blocking=True)
    
    # This is a synchronization point
    # It will wait for previous two lines
    output = model(features)

8. 使用 PyTorch JIT 將點操作融合到單個 kernel 中

點操作包括常見的數學操作，通常是內存受限的。PyTorch JIT 會自動將相鄰的點操作融合到一個內核中，以保存多次內存讀 / 寫操作。例如，通過將 5 個核融合成 1 個核，gelu函數可以被加速 4 倍。

@torch.jit.script # JIT decorator
def fused_gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

9 & 10. 在使用混合精度的 FP16 時，對於所有不同架構設計，設置圖像尺寸和 batch size 爲 8 的倍數

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

這些實驗顯示設置輸出維度和 batch size 大小爲 8 的倍數，比如 (33712、4088、4096) 相比 33708，batch size 爲 4084 或者 4095 這些不能被 8 整除的數可以加速計算 1.3 倍到 4 倍。加速度大小取決於過程類型 (例如，向前傳遞或梯度計算) 和 cuBLAS 版本。特別是，如果你使用 NLP，請記住檢查輸出維度，這通常是詞彙表大小。

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

除了將 batch size 大小設置爲 8 的倍數外，我們還將 batch size 大小最大化，直到它達到 GPU 的內存限制。這樣，我們可以用更少的時間來完成一個 epoch。

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

有些操作不需要 float64 或 float32 的精度。因此，將操作設置爲較低的精度可以節省內存和執行時間。對於各種應用，英偉達報告稱具有 Tensor Cores 的 GPU 的混合精度可以提高 3.5 到 25 倍的速度。

值得注意的是，通常矩陣越大，混合精度加速度越高。在較大的神經網絡中 (例如 BERT)，實驗表明混合精度可以加快 2.75 倍的訓練，並減少 37% 的內存使用。具有 Volta, Turing, Ampere 或 Hopper 架構的較新的 GPU 設備(例如，T4, V100, RTX 2060, 2070, 2080, 2080 Ti, A100, RTX 3090, RTX 3080，和 RTX 3070) 可以從混合精度中受益更多，因爲他們有 Tensor Core 架構，它相比 CUDA cores 有特殊的優化。

帶有 Tensor Core 的 NVIDIA 架構支持不同的精度

值得一提的是，採用 Hopper 架構的 H100 預計將於 2022 年第三季度發佈，支持 FP8 (float8)。PyTorch AMP 可能會支持 FP8(目前 v1.11.0 還不支持 FP8)。

在實踐中，你需要在模型精度性能和速度性能之間找到一個最佳點。我之前確實發現混合精度可能會降低模型的精度，這取決於算法，數據和問題。

使用自動混合精度 (AMP) 很容易在 PyTorch 中利用混合精度。PyTorch 中的默認浮點類型是 float32。AMP 將通過使用 float16 來進行一組操作(例如，matmul, linear, conv2d) 來節省內存和時間。AMP 會自動 cast 到 float32 的一些操作 (例如，mse_loss, softmax等)。有些操作 (例如add) 可以操作最寬的輸入類型。例如，如果一個變量是 float32，另一個變量是 float16，那麼加法結果將是 float32。

autocast自動應用精度到不同的操作。因爲損失和梯度是按照 float16 精度計算的，當它們太小時，梯度可能會 “下溢” 並變成零。GradScaler通過將損失乘以一個比例因子來防止下溢，根據比例損失計算梯度，然後在優化器更新權重之前取消梯度的比例。如果縮放因子太大或太小，並導致inf或NaN，則縮放因子將在下一個迭代中更新縮放因子。

scaler = GradScaler()
for features, target in data:
    # Forward pass with mixed precision
    with torch.cuda.amp.autocast(): # autocast as a context manager
        output = model(features)
        loss = criterion(output, target)    
    
    # Backward pass without mixed precision
    # It's not recommended to use mixed precision for backward pass
    # Because we need more precise loss
    scaler.scale(loss).backward()    
    
    # scaler.step() first unscales the gradients .
    # If these gradients contain infs or NaNs, 
    # optimizer.step() is skipped.
    scaler.step(optimizer)     
    
    # If optimizer.step() was skipped,
    # scaling factor is reduced by the backoff_factor in GradScaler()
    scaler.update()

你也可以使用autocast 作爲前向傳遞函數的裝飾器。

class AutocastModel(nn.Module):
    ...
    @autocast() # autocast as a decorator
    def forward(self, input):
        x = self.model(input)
        return x

12. 在優化器更新權重之前將梯度設置爲 None

通過model.zero_grad()或optimizer.zero_grad()將對所有參數執行memset ，並通過讀寫操作更新梯度。但是，將梯度設置爲None將不會執行memset，並且將使用 “只寫” 操作更新梯度。因此，設置梯度爲None更快。

# Reset gradients before each step of optimizer
for param in model.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 個 batch 再更新梯度，模擬大 batch size

這個技巧是關於從更多的數據樣本積累梯度，以便對梯度的估計更準確，權重更新更接近局部 / 全局最小值。這在 batch size 較小的情況下更有幫助 (由於 GPU 內存限制較小或每個樣本的數據量較大)。

for i, (features, target) in enumerate(dataloader):
    # Forward pass
    output = model(features)
    loss = criterion(output, target)    
    
    # Backward pass
    loss.backward()    
    
    # Only update weights every other 2 iterations
    # Effective batch size is doubled
    if (i+1) % 2 == 0 or (i+1) == len(dataloader):
        # Update weights
        optimizer.step()        
        # Reset the gradients to None
        optimizer.zero_grad(set_to_none=True)

14. 在推理和驗證的時候禁用梯度計算

實際上，如果只計算模型的輸出，那麼梯度計算對於推斷和驗證步驟並不是必需的。PyTorch 使用一箇中間內存緩衝區來處理requires_grad=True變量中涉及的操作。因此，如果我們知道不需要任何涉及梯度的操作，通過禁用梯度計算來進行推斷 / 驗證，就可以避免使用額外的資源。

# torch.no_grad() as a context manager:
    with torch.no_grad():
    output = model(input)
    
# torch.no_grad() as a function decorator:
@torch.no_grad()
def validation(model, input):
    output = model(input)
return output

15. torch.backends.cudnn.benchmark = True

在訓練循環之前設置torch.backends.cudnn.benchmark = True可以加速計算。由於計算不同內核大小卷積的 cuDNN 算法的性能不同，自動調優器可以運行一個基準來找到最佳算法。當你的輸入大小不經常改變時，建議開啓這個設置。如果輸入大小經常改變，那麼自動調優器就需要太頻繁地進行基準測試，這可能會損害性能。它可以將向前和向後傳播速度提高 1.27x 到 1.70x。

torch.backends.cudnn.benchmark = True

16. 對於 4D NCHW Tensors 使用通道在最後的內存格式

4D NCHW 重新組織成 NHWC 格式

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

目前，它仍處於 beta 測試階段，僅支持 4D NCHW 張量和一組模型 (例如，alexnet，mnasnet家族，mobilenet_v2，resnet家族，shufflenet_v2，squeezenet1，vgg家族)。但我可以肯定，這將成爲一個標準的優化。

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

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

# Stride is the gap between one element to the next one 
# in a dimension.
print(x.stride()) 

# (3072, 1024, 32, 1)# Convert the tensor to NHWC in memory
x2 = x.to(memory_format=torch.channels_last)

print(x2.shape)  # (10, 3, 32, 32) as dimensions order preserved
print(x2.stride())  # (3072, 1, 96, 3), which are smaller
print((x==x2).all()) # True because the values were not changed

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

這是可行的，因爲在數學上，bias 可以通過 batch normalization 的均值減法來抵消。我們可以節省模型參數、運行時的內存。

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

18. 使用 DistributedDataParallel代替DataParallel

對於多 GPU 來說，即使只有單個節點，也總是優先使用 DistributedDataParallel而不是 DataParallel ，因爲 DistributedDataParallel 應用於多進程，併爲每個 GPU 創建一個進程，從而繞過 Python 全局解釋器鎖 (GIL) 並提高速度。

總結

在這篇文章中，我列出了一個清單，並提供了 18 個 PyTorch 技巧的代碼片段。然後，我逐一解釋了它們在不同方面的工作原理和原因，包括數據加載、數據操作、模型架構**、**訓練、推斷、cnn 特定的優化和分佈式計算。一旦你深入理解了它們的工作原理，你可能會找到適用於任何深度學習框架中的深度學習建模的通用原則。

參考資料

英文原文：https://towardsdatascience.com/optimize-pytorch-performance-for-speed-and-memory-efficiency-2022-84f453916ea6

煉丹筆記 一本有仙氣的筆記，記錄了 AI 算法圈裏的祕密…

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/XbNo22gVmAjEndzLLmnmxQ