为输入数据确保正确的内存格式，可以显著影响 PyTorch 视觉模型的运行时间。如果拿不准，请选择“通道最后”（Channels Last）内存格式。
当在 PyTorch 中处理接收多媒体（例如图像张量）作为输入的视觉模型时，张量的内存格式会显著影响在使用 CPU 后端配合 XNNPACK 时，模型在移动平台上的推理执行速度。这一点对于服务器平台上的训练和推理同样适用，但延迟对于移动设备和用户而言尤为关键。
本文大纲
- 深入探讨 C++ 中的矩阵存储/内存表示。介绍行优先和列优先顺序。
- 遍历矩阵时的顺序与存储表示相同或不同所产生的影响，以及相关示例。
- Cachegrind 简介；一个用于检查代码缓存友好性的工具。
- PyTorch 算子支持的内存格式。
- 确保 XNNPACK 优化下模型高效执行的最佳实践示例。
C++ 中的矩阵存储表示
图像作为多维张量被馈送到 PyTorch 机器学习模型中。这些张量具有特定的内存格式。为了更好地理解这一概念，让我们看看二维矩阵在内存中是如何存储的。
广义上讲，在内存中高效存储多维数据主要有两种方式。
- 行优先顺序 (Row Major Order)： 在这种格式中，矩阵按行顺序存储，每一行在内存中都存储在下一行之前。即第 N 行位于第 N+1 行之前。
- 列优先顺序 (Column Major Order)： 在这种格式中，矩阵按列顺序存储，每一列在内存中都存储在下一列之前。即第 N 列位于第 N+1 列之前。
您可以在下方以图形方式查看二者的差异。
C++ 以行优先格式存储多维数据。
高效访问二维矩阵元素
与存储格式类似，访问二维矩阵中的数据也有两种方式。
- 先遍历行： 在处理下一行的任何元素之前，先处理当前行的所有元素。
- 先遍历列： 在处理下一列的任何元素之前，先处理当前列的所有元素。
为了获得最大效率，应始终以与数据存储格式相同的顺序访问数据。即，如果数据以行优先顺序存储，则应尝试以该顺序访问它。
下面的代码 (main.cpp) 展示了访问 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() 导致的 L1 数据缓存未命中次数比 loop1() 多得多（约 16 倍）。这就是为什么 loop1() 比 loop2() 快约 15 倍的原因。
PyTorch 算子支持的内存格式
虽然 PyTorch 算子期望所有张量都采用通道优先 (NCHW) 维度格式，但 PyTorch 算子支持 3 种输出内存格式。
- 连续 (Contiguous)： 张量内存与张量的维度顺序一致。
- 通道最后 (ChannelsLast)： 无论维度顺序如何，二维（图像）张量在内存中都被布局为 HWC 或 NHWC（N：批次，H：高度，W：宽度，C：通道）张量。维度可以以任何顺序排列。
- 通道最后 3D (ChannelsLast3d)： 对于三维张量（视频张量），内存布局为 THWC（时间、高度、宽度、通道）或 NTHWC（N：批次，T：时间，H：高度，W：宽度，C：通道）格式。维度可以以任何顺序排列。
视觉模型之所以偏好 ChannelsLast，是因为 PyTorch 使用的 XNNPACK（内核加速库）期望所有输入都采用 通道最后 格式。因此，如果模型的输入不是通道最后格式，则必须首先将其转换为该格式，这会产生额外的开销。
此外，大多数 PyTorch 算子会保留输入张量的内存格式，因此如果输入是通道优先格式，算子则需要先将其转换为通道最后，执行操作，然后再转回通道优先。
当考虑到加速算子在通道最后内存格式下表现更好这一事实时，您会发现让算子返回通道最后格式对后续的算子调用更有利，否则每个算子最终都要执行转换为通道最后的操作（如果这对该特定算子更高效的话）。
引用自 XNNPACK 主页：
“XNNPACK 中的所有算子都支持 NHWC 布局，此外还允许在通道维度上使用自定义步长 (stride)”。
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)
对于连续和通道最后格式，输入形状保持不变。但在内部，张量的布局如步长 (strides) 所示发生了变化。现在，跨越通道所需的跳跃次数仅为 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 模型的性能。