概述
在运行视觉模型时，内存格式对性能有显著影响，通常 Channels Last 由于更好的数据局部性，在性能方面更具优势。
本博客将介绍内存格式的基本概念，并演示在 Intel® Xeon® 可扩展处理器上，使用 Channels Last 在流行 PyTorch 视觉模型上的性能优势。
内存格式介绍
内存格式指的是数据表示，它描述了多维 (nD) 数组在线性 (1D) 内存地址空间中的存储方式。内存格式的概念有两个方面：
- 物理顺序 是数据在物理内存中的存储布局。对于视觉模型，我们通常讨论 NCHW、NHWC。这些是物理内存布局的描述，也分别称为 Channels First 和 Channels Last。
- 逻辑顺序 是描述张量形状和步幅的约定。在 PyTorch 中，此约定为 NCHW。无论物理顺序如何，张量形状和步幅将始终以 NCHW 的顺序表示。
图 1 显示了形状为 [1, 3, 4, 4] 的张量在 Channels First 和 Channels Last 内存格式下的物理内存布局（通道分别表示为 R、G、B）
图 1 Channels First 和 Channels Last 的物理内存布局
内存格式传播
PyTorch 内存格式传播的通用规则是保留输入张量的内存格式。这意味着 Channels First 输入将生成 Channels First 输出，而 Channels Last 输入将生成 Channels Last 输出。
对于卷积层，PyTorch 默认使用 oneDNN (oneAPI 深度神经网络库) 在 Intel CPU 上实现最佳性能。由于直接使用 Channels First 内存格式物理上无法实现高度优化的性能，因此输入和权重首先转换为块状格式，然后进行计算。oneDNN 可能会根据输入形状、数据类型和硬件架构选择不同的块状格式，以实现向量化和缓存重用。块状格式对 PyTorch 是不透明的，因此输出需要转换回 Channels First。尽管块状格式可以带来最佳计算性能，但格式转换可能会增加开销，从而抵消性能增益。
另一方面，oneDNN 已针对 Channels Last 内存格式进行优化，以直接使用它来实现最佳性能，PyTorch 将简单地将内存视图传递给 oneDNN。这意味着节省了输入和输出张量的转换。图 2 显示了 PyTorch CPU 上卷积的内存格式传播行为（实心箭头表示内存格式转换，虚线箭头表示内存视图）
图 2 CPU 卷积内存格式传播
在 PyTorch 中，默认的内存格式是 Channels First。如果某个特定运算符不支持 Channels Last，则 NHWC 输入将被视为不连续的 NCHW，因此会回退到 Channels First，这将消耗 CPU 上先前的内存带宽并导致次优性能。
因此，扩大 Channels Last 的支持范围以获得最佳性能非常重要。我们已经为 CV 领域常用的运算符实现了 Channels Last 内核，适用于推理和训练，例如：
- 激活函数（例如，ReLU、PReLU 等）
- 卷积（例如，Conv2d）
- 归一化（例如，BatchNorm2d、GroupNorm 等）
- 池化（例如，AdaptiveAvgPool2d、MaxPool2d 等）
- 混洗（例如，ChannelShuffle、PixelShuffle）
有关详细信息，请参阅 Operators-with-Channels-Last-support。
Channels Last 的原生级优化
如上所述，PyTorch 使用 oneDNN 在 Intel CPU 上为卷积实现最佳性能。其余内存格式感知运算符在 PyTorch 原生级别进行优化，不需要任何第三方库支持。
- 缓存友好的并行化方案： 对所有内存格式感知运算符保持相同的并行化方案，这将有助于在将每个层的输出传递给下一层时增加数据局部性。
- 多架构向量化： 通常，我们可以在 Channels Last 内存格式上对最内层维度进行向量化。并且每个向量化的 CPU 内核都将为 AVX2 和 AVX512 生成。
在贡献 Channels Last 内核时，我们也尽力优化了 Channels First 对应的内核。事实是，某些运算符在 Channels First 上物理上无法实现最佳性能，例如卷积、池化等。
在 Channels Last 上运行视觉模型
Channels Last 相关的 API 已在 PyTorch 内存格式教程 中有详细说明。通常，我们可以通过以下方式将 4D 张量从 Channels First 转换为 Channels Last：
# convert x to channels last
# suppose x’s shape is (N, C, H, W)
# then x’s stride will be (HWC, 1, WC, C)
x = x.to(memory_format=torch.channels_last)
要在 Channels Last 内存格式上运行模型，只需将输入和模型转换为 Channels Last 即可。以下是一个最小示例，展示如何在 Channels Last 内存格式下运行带 TorchVision 的 ResNet50：
import torch
from torchvision.models import resnet50
N, C, H, W = 1, 3, 224, 224
x = torch.rand(N, C, H, W)
model = resnet50()
model.eval()
# convert input and model to channels last
x = x.to(memory_format=torch.channels_last)
model = model.to(memory_format=torch.channels_last)
model(x)
Channels Last 优化是在原生内核级别实现的，这意味着您也可以将其他功能（例如 torch.fx 和 torch script）与 Channels Last 一起应用。
性能提升
我们在 Intel® Xeon® Platinum 8380 CPU @ 2.3 GHz 上对 TorchVision 模型的推理性能进行了基准测试，每个插槽一个实例（批处理大小 = 2 x 物理核心数）。结果显示 Channels Last 比 Channels First 有 1.3 倍到 1.8 倍的性能提升。
性能提升主要来自两个方面：
- 对于卷积层，Channels Last 节省了激活函数转换为块状格式的内存格式转换，从而提高了整体计算效率。
- 对于池化和上采样层，Channels Last 可以沿着最内层维度（例如“C”）使用向量化逻辑，而 Channels First 则不能。
对于内存格式不感知的层，Channels Last 和 Channels First 具有相同的性能。
结论与未来工作
在本博客中，我们介绍了 Channels Last 的基本概念，并演示了在视觉模型上使用 Channels Last 在 CPU 上的性能优势。当前的工作仅限于 2D 模型，我们将在不久的将来将优化工作扩展到 3D 模型！
致谢
本博客中呈现的结果是 Meta 和 Intel PyTorch 团队的共同努力。特别感谢 Meta 的 Vitaly Fedyunin 和 Wei Wei，他们付出了宝贵的时间并提供了实质性帮助！我们共同在改进 PyTorch CPU 生态系统的道路上又迈出了一步。