高效 PyTorch：張量記憶體格式很重要

確保輸入使用正確的記憶體格式可以顯著影響 PyTorch 視覺模型的執行時間。如有疑問，請選擇 Channels Last 記憶體格式。

在處理接受多媒體（例如影像張量）作為輸入的 PyTorch 視覺模型時，張量的記憶體格式可以顯著影響 在使用 CPU 後端和 XNNPACK 的移動平臺上模型的推理執行速度。這對於伺服器平臺上的訓練和推理也適用，但對於移動裝置和使用者而言，延遲尤為關鍵。

本文大綱

深入探討 C++ 中的矩陣儲存/記憶體表示。介紹行主序和列主序。
按照與儲存表示相同或不同的順序遍歷矩陣的影響，並附帶一個示例。
Cachegrind 簡介；一個檢查程式碼快取友好性的工具。
PyTorch 運算子支援的記憶體格式。
確保使用 XNNPACK 最佳化高效執行模型的最佳實踐示例。

C++ 中的矩陣儲存表示

影像以多維張量的形式輸入到 PyTorch ML 模型中。這些張量具有特定的記憶體格式。為了更好地理解這個概念，讓我們看看二維矩陣是如何儲存在記憶體中的。

廣義上講，有 2 種主要方式可以高效地在記憶體中儲存多維資料。

行主序： 在這種格式中，矩陣按行順序儲存，每行在記憶體中儲存在下一行之前。即，第 N 行在第 N+1 行之前。
列主序： 在這種格式中，矩陣按列順序儲存，每列在記憶體中儲存在下一列之前。即，第 N 列在第 N+1 列之前。

您可以在下面的圖表中看到差異。

C++ stores multi-dimensional data in row-major format.
C++ 以行主序格式儲存多維資料。

高效訪問二維矩陣的元素

與儲存格式類似，有 2 種方法可以訪問二維矩陣中的資料。

首先遍歷行： 在處理下一行的任何元素之前，處理一行的所有元素。
首先遍歷列： 在處理下一列的任何元素之前，處理一列的所有元素。

為了達到最大效率，應始終以與資料儲存格式相同的格式訪問資料。即，如果資料以行主序儲存，則應嘗試以該順序訪問它。

下面的程式碼 (main.cpp) 展示了 2 種方法訪問 4000×4000 二維矩陣的所有元素。

#include <iostream>
#include <chrono>

// loop1 accesses data in matrix 'a' in row major order,
// since i is the outer loop variable, and j is the
// inner loop variable.
int loop1(int a[4000][4000]) {
 int s = 0;
 for (int i = 0; i < 4000; ++i) {
   for (int j = 0; j < 4000; ++j) {
     s += a[i][j];
   }
 }
 return s;
}

// loop2 accesses data in matrix 'a' in column major order
// since j is the outer loop variable, and i is the
// inner loop variable.
int loop2(int a[4000][4000]) {
 int s = 0;
 for (int j = 0; j < 4000; ++j) {
   for (int i = 0; i < 4000; ++i) {
     s += a[i][j];
   }
 }
 return s;
}

int main() {
 static int a[4000][4000] = {0};
 for (int i = 0; i < 100; ++i) {
   int x = rand() % 4000;
   int y = rand() % 4000;
   a[x][y] = rand() % 1000;
 }

 auto start = std::chrono::high_resolution_clock::now();
 auto end = start;
 int s = 0;

#if defined RUN_LOOP1
 start = std::chrono::high_resolution_clock::now();

 s = 0;
 for (int i = 0; i < 10; ++i) {
   s += loop1(a);
   s = s % 100;
 }
 end = std::chrono::high_resolution_clock::now();

 std::cout << "s = " << s << std::endl;
 std::cout << "Time for loop1: "
   << std::chrono::duration<double, std::milli>(end - start).count()
   << "ms" << std::endl;
#endif

#if defined RUN_LOOP2
 start = std::chrono::high_resolution_clock::now();
 s = 0;
 for (int i = 0; i < 10; ++i) {
   s += loop2(a);
   s = s % 100;
 }
 end = std::chrono::high_resolution_clock::now();

 std::cout << "s = " << s << std::endl;
 std::cout << "Time for loop2: "
   << std::chrono::duration<double, std::milli>(end - start).count()
   << "ms" << std::endl;
#endif
}


Let’s build and run this program and see what it prints.

g++ -O2 main.cpp -DRUN_LOOP1 -DRUN_LOOP2
./a.out


Prints the following:

s = 70
Time for loop1: 77.0687ms
s = 70
Time for loop2: 1219.49ms

loop1() 比 loop2() 快 15 倍。為什麼會這樣？讓我們在下面找出答案！

使用 Cachegrind 測量快取未命中

Cachegrind 是一個快取分析工具，用於檢視程式導致了多少 I1（一級指令）、D1（一級資料）和 LL（末級）快取未命中。

讓我們只用 loop1() 和只用 loop2() 構建程式，看看這些函式各自的快取友好性。

構建並執行/分析只包含 loop1() 的程式

g++ -O2 main.cpp -DRUN_LOOP1
valgrind --tool=cachegrind ./a.out

列印：

==3299700==
==3299700== I   refs:      643,156,721
==3299700== I1  misses:          2,077
==3299700== LLi misses:          2,021
==3299700== I1  miss rate:        0.00%
==3299700== LLi miss rate:        0.00%
==3299700==
==3299700== D   refs:      160,952,192  (160,695,444 rd   + 256,748 wr)
==3299700== D1  misses:     10,021,300  ( 10,018,723 rd   +   2,577 wr)
==3299700== LLd misses:     10,010,916  ( 10,009,147 rd   +   1,769 wr)
==3299700== D1  miss rate:         6.2% (        6.2%     +     1.0%  )
==3299700== LLd miss rate:         6.2% (        6.2%     +     0.7%  )
==3299700==
==3299700== LL refs:        10,023,377  ( 10,020,800 rd   +   2,577 wr)
==3299700== LL misses:      10,012,937  ( 10,011,168 rd   +   1,769 wr)
==3299700== LL miss rate:          1.2% (        1.2%     +     0.7%  )

構建並執行/分析只包含 loop2() 的程式

g++ -O2 main.cpp -DRUN_LOOP2
valgrind --tool=cachegrind ./a.out

列印：

==3300389==
==3300389== I   refs:      643,156,726
==3300389== I1  misses:          2,075
==3300389== LLi misses:          2,018
==3300389== I1  miss rate:        0.00%
==3300389== LLi miss rate:        0.00%
==3300389==
==3300389== D   refs:      160,952,196  (160,695,447 rd   + 256,749 wr)
==3300389== D1  misses:    160,021,290  (160,018,713 rd   +   2,577 wr)
==3300389== LLd misses:     10,014,907  ( 10,013,138 rd   +   1,769 wr)
==3300389== D1  miss rate:        99.4% (       99.6%     +     1.0%  )
==3300389== LLd miss rate:         6.2% (        6.2%     +     0.7%  )
==3300389==
==3300389== LL refs:       160,023,365  (160,020,788 rd   +   2,577 wr)
==3300389== LL misses:      10,016,925  ( 10,015,156 rd   +   1,769 wr)
==3300389== LL miss rate:          1.2% (        1.2%     +     0.7%  )

兩次執行的主要區別是

D1 未命中： 10M 對比 160M
D1 未命中率： 6.2% 對比 99.4%

正如您所看到的，loop2() 導致了比 loop1() 多得多的（大約多 16 倍）L1 資料快取未命中。這就是為什麼loop1()比 loop2() 快約 15 倍。

PyTorch 運算子支援的記憶體格式

雖然 PyTorch 運算子期望所有張量都採用 Channels First (NCHW) 維度格式，但 PyTorch 運算子支援 3 種輸出記憶體格式。

連續： 張量記憶體的順序與張量維度相同。
ChannelsLast： 無論維度順序如何，2D（影像）張量在記憶體中均以 HWC 或 NHWC （N：批次，H：高度，W：寬度，C：通道）張量佈局。維度可以以任何順序排列。
ChannelsLast3d： 對於 3D 張量（影片張量），記憶體以 THWC（時間、高度、寬度、通道）或 NTHWC（N：批次，T：時間，H：高度，W：寬度，C：通道）格式佈局。維度可以以任何順序排列。

視覺模型首選 ChannelsLast 的原因是，PyTorch 使用的 XNNPACK （核心加速庫）期望所有輸入都採用 Channels Last 格式，因此如果模型的輸入不是 Channels Last，則必須先將其轉換為 Channels Last，這是一個額外的操作。

此外，大多數 PyTorch 運算子會保留輸入張量的記憶體格式，因此如果輸入是 Channels First，則運算子需要先轉換為 Channels Last，然後執行操作，然後再轉換回 Channels First。

當您將其與加速運算子在 Channels Last 記憶體格式下工作得更好這一事實結合起來時，您會發現讓運算子返回 Channels Last 記憶體格式對於後續的運算子呼叫更好，否則您最終會使每個運算子都轉換為 Channels Last（如果這對特定運算子更有效）。

來自 XNNPACK 主頁

“XNNPACK 中的所有運算子都支援 NHWC 佈局，但此外還允許沿通道維度自定義步幅。”

PyTorch 最佳實踐

要從 PyTorch 視覺模型中獲得最佳效能，最好的方法是確保您的輸入張量在饋入模型之前採用 Channels Last 記憶體格式。

透過最佳化模型以使用 XNNPACK 後端（只需在您的 torchscripted 模型上呼叫 optimize_for_mobile()），您可以獲得更快的速度。請注意，如果輸入是連續的，XNNPACK 模型將執行得更慢，因此務必確保它採用 Channels-Last 格式。

顯示加速的工作示例

在 Google Colab 上執行此示例 – 請注意，colab CPU 上的執行時可能無法反映準確的效能；建議在您的本地機器上執行此程式碼。

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile
import torch.backends.xnnpack
import time

print("XNNPACK is enabled: ", torch.backends.xnnpack.enabled, "\n")

N, C, H, W = 1, 3, 200, 200
x = torch.rand(N, C, H, W)
print("Contiguous shape: ", x.shape)
print("Contiguous stride: ", x.stride())
print()

xcl = x.to(memory_format=torch.channels_last)
print("Channels-Last shape: ", xcl.shape)
print("Channels-Last stride: ", xcl.stride())

## Outputs:
 
# XNNPACK is enabled:  True
 
# Contiguous shape:  torch.Size([1, 3, 200, 200])
# Contiguous stride:  (120000, 40000, 200, 1)
 
# Channels-Last shape:  torch.Size([1, 3, 200, 200])
# Channels-Last stride:  (120000, 1, 600, 3)

對於連續和 Channels Last 格式，輸入形狀保持不變。然而，在內部，張量的佈局已經改變，正如您在步幅中看到的那樣。現在，跨通道所需的跳轉次數僅為 1（而不是連續張量中的 40000）。這種更好的資料區域性性意味著卷積層可以更快地訪問給定畫素的所有通道。現在讓我們看看記憶體格式如何影響執行時

from torchvision.models import resnet34, resnet50, resnet101

m = resnet34(pretrained=False)
# m = resnet50(pretrained=False)
# m = resnet101(pretrained=False)

def get_optimized_model(mm):
  mm = mm.eval()
  scripted = torch.jit.script(mm)
  optimized = optimize_for_mobile(scripted)  # explicitly call the xnnpack rewrite 
  return scripted, optimized


def compare_contiguous_CL(mm):
  # inference on contiguous
  start = time.perf_counter()
  for i in range(20):
    mm(x)
  end = time.perf_counter()
  print("Contiguous: ", end-start)

  # inference on channels-last
  start = time.perf_counter()
  for i in range(20):
    mm(xcl)
  end = time.perf_counter()
  print("Channels-Last: ", end-start)

with torch.inference_mode():
  scripted, optimized = get_optimized_model(m)

  print("Runtimes for torchscripted model: ")
  compare_contiguous_CL(scripted.eval())
  print()
  print("Runtimes for mobile-optimized model: ")
  compare_contiguous_CL(optimized.eval())

   
## Outputs (on an Intel Core i9 CPU):
 
# Runtimes for torchscripted model:
# Contiguous:  1.6711160129999598
# Channels-Last:  1.6678222839999535
 
# Runtimes for mobile-optimized model:
# Contiguous:  0.5712863490000473
# Channels-Last:  0.46113000699995155

結論

輸入張量的記憶體佈局可以顯著影響模型的執行時間。對於視覺模型，請優先選擇 Channels Last 記憶體格式，以充分利用您的 PyTorch 模型。

高效 PyTorch：張量記憶體格式至關重要