跳转到主要内容
博客

高效 PyTorch：张量内存格式至关重要

作者： 2021 年 12 月 15 日2024 年 11 月 15 日暂无评论

确保为输入选择正确的内存格式可以显著影响 PyTorch 视觉模型的运行时间。如有疑问，请选择通道优先（Channels Last）内存格式。

在 PyTorch 中处理接受多媒体（例如图像张量）作为输入的视觉模型时，张量的内存格式可以显著影响**在移动平台上使用 CPU 后端和 XNNPACK 时模型的推理执行速度**。这对于服务器平台上的训练和推理也适用，但延迟对于移动设备和用户尤其关键。

本文大纲

  1. 深入探讨 C++ 中的矩阵存储/内存表示。介绍行主序和列主序。
  2. 以与存储表示相同或不同顺序遍历矩阵的影响，并附带示例。
  3. 介绍 Cachegrind；一个用于检查代码缓存友好性的工具。
  4. PyTorch 运算符支持的内存格式。
  5. 确保使用 XNNPACK 优化高效执行模型的最佳实践示例

C++ 中的矩阵存储表示

图像以多维张量形式输入到 PyTorch 机器学习模型中。这些张量具有特定的内存格式。为了更好地理解这个概念，让我们看看 2 维矩阵如何存储在内存中。

广义地说，在内存中高效存储多维数据有两种主要方式。

  1. **行主序：**在这种格式中，矩阵按行存储，每行在内存中存储在下一行之前。即，第 N 行位于第 N+1 行之前。
  2. **列主序：**在这种格式中，矩阵按列存储，每列在内存中存储在下一列之前。即，第 N 列位于第 N+1 列之前。

您可以在下方图形中查看差异。

C++ stores multi-dimensional data in row-major format.
C++ 以行主序格式存储多维数据。

高效访问 2 维矩阵的元素

类似于存储格式，有两种方式访问 2 维矩阵中的数据。

  1. **先按行遍历：**在处理下一行的任何元素之前，先处理一行的所有元素。
  2. **先按列遍历：**在处理下一列的任何元素之前，先处理一列的所有元素。

为了最大限度地提高效率，应始终以与数据存储格式相同的格式访问数据。也就是说，如果数据以行主序存储，那么应尝试以该顺序访问它。

下面的代码 (main.cpp) 显示了访问 2 维 4000x4000 矩阵所有元素的两种方式。

#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%  )

两次运行之间的主要区别是

  1. **D1 未命中：**10M 对 160M
  2. **D1 未命中率：**6.2% 对 99.4%

如您所见，`loop2()` 导致的一级数据缓存未命中比 loop1() 多得多（**约 16 倍**）。这就是为什么 `loop1()` 比 loop2() 快约 15 倍的原因。

PyTorch 运算符支持的内存格式

虽然 PyTorch 运算符期望所有张量都采用通道优先 (NCHW) 维度格式，但 PyTorch 运算符支持 3 种输出内存格式。

  1. **连续 (Contiguous)：**张量内存在物理上与张量的维度顺序一致。
  2. **通道优先 (ChannelsLast)：**无论维度顺序如何，2D（图像）张量在内存中均按 HWC 或NHWC（N：批次，H：高度，W：宽度，C：通道）张量排列。维度可以按任意顺序排列。
  3. **通道优先 3D (ChannelsLast3d)：**对于 3D 张量（视频张量），内存按 THWC（时间、高度、宽度、通道）或 NTHWC（N：批次，T：时间，H：高度，W：宽度，C：通道）格式排列。维度可以按任意顺序排列。

通道优先 (ChannelsLast) 之所以受视觉模型青睐，是因为 PyTorch 使用的 XNNPACK（内核加速库）要求所有输入都采用**通道优先**格式，因此如果模型的输入不是通道优先，则必须先将其转换为通道优先，这会增加额外的操作。

此外，大多数 PyTorch 运算符会保留输入张量的内存格式，因此如果输入是通道优先，则运算符需要首先转换为通道优先，然后执行操作，然后再转换回通道优先。

当您结合加速运算符在通道优先内存格式下表现更好的事实时，您会发现让运算符返回通道优先内存格式对于后续的运算符调用更好，否则您最终会导致每个运算符都转换为通道优先（如果这对特定运算符更有效）。

摘自 XNNPACK 主页

“XNNPACK 中的所有运算符都支持 NHWC 布局，但额外允许沿通道维度的自定义步幅。”

PyTorch 最佳实践

要让 PyTorch 视觉模型发挥最大性能，最佳方法是确保在将输入张量输入模型之前，其处于**通道优先**内存格式。

您可以通过优化模型以使用 XNNPACK 后端（只需在 torchscript 模型上调用 `optimize_for_mobile()`）来获得更快的速度。请注意，如果输入是连续的，XNNPACK 模型将运行缓慢，因此务必确保它采用通道优先格式。

显示加速的示例

在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)

对于连续和通道优先格式，输入形状保持不变。然而，张量的内部布局已更改，如您在步长中所示。现在，跨通道所需的跳跃次数仅为 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

结论

输入张量的内存布局可以显著影响模型的运行时间。对于视觉模型，优先选择**通道优先**内存格式，以充分利用您的 PyTorch 模型。

参考文献