确保您的输入采用正确的内存格式可以显著影响 PyTorch 视觉模型的运行时间。如有疑问，请选择 Channels Last 内存格式。
当在 PyTorch 中处理接受多媒体（例如图像张量）作为输入的视觉模型时，张量的内存格式会显著影响在使用 CPU 后端和 XNNPACK 时，模型在移动平台上的推理执行速度。这对于服务器平台上的训练和推理也同样适用，但延迟对于移动设备和用户而言尤为关键。
本文大纲
- 深入探讨 C++ 中的矩阵存储/内存表示。介绍 行优先和列优先顺序。
- 以与存储表示相同或不同的顺序循环遍历矩阵的影响，并提供示例。
- Cachegrind 简介；一个用于检查代码缓存友好程度的工具。
- PyTorch 运算符支持的内存格式。
- 最佳实践示例，以确保通过 XNNPACK 优化实现高效的模型执行
C++ 中的矩阵存储表示
图像以多维张量的形式馈送到 PyTorch ML 模型中。这些张量具有特定的内存格式。为了更好地理解这个概念，让我们看看如何将 2 维矩阵存储在内存中。
广义上讲，有效存储内存中的多维数据主要有两种方式。
- 行优先顺序： 在此格式中，矩阵以行顺序存储，每行在内存中都存储在下一行之前。即，行 N 在行 N+1 之前。
- 列优先顺序： 在此格式中，矩阵以列顺序存储，每列在内存中都存储在下一列之前。即，列 N 在列 N+1 之前。
您可以在下方以图形方式查看差异。
C++ 以行优先格式存储多维数据。
高效访问 2 维矩阵的元素
与存储格式类似，访问 2 维矩阵中的数据有两种方式。
- 先循环遍历行： 在处理下一行的任何元素之前，先处理完一行的所有元素。
- 先循环遍历列： 在处理下一列的任何元素之前，先处理完一列的所有元素。
为了获得最大效率，应始终以数据存储的格式访问数据。即，如果数据以行优先顺序存储，则应尝试以该顺序访问它。
下面的代码 (main.cpp) 展示了 2 种方式 访问 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% )
两次运行之间的主要区别在于
- D1 未命中： 1 千万 对比 1.6 亿
- D1 未命中率： 6.2% 对比 99.4%
如您所见，loop2() 比 loop1() 导致多得多的（约 16 倍）L1 数据缓存未命中。这就是 loop1() 比 loop2() 快约 15 倍的原因。
PyTorch 运算符支持的内存格式
虽然 PyTorch 运算符期望所有张量都采用 Channels First (NCHW) 维度格式，但 PyTorch 运算符支持 3 种输出 内存格式。
- Contiguous（连续）： 张量内存与张量的维度顺序相同。
- ChannelsLast（通道优先）： 无论维度顺序如何，2 维（图像）张量在内存中都布局为 HWC 或 NHWC（N：批次，H：高度，W：宽度，C：通道）张量。维度可以按任何顺序排列。
- ChannelsLast3d（3D 通道优先）： 对于 3 维张量（视频张量），内存以 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 后端（只需在您的 torchscript 模型上调用 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 模型。
