概述
内存格式对运行视觉模型的性能有显著影响。从性能角度来看，由于具有更好的数据局部性，Channels Last（通道在后）格式通常更具优势。
本博客将介绍内存格式的基本概念，并展示在英特尔® 至强®（Intel® Xeon®）可扩展处理器上，使用 Channels Last 对主流 PyTorch 视觉模型带来的性能提升。
内存格式简介
内存格式是指描述多维 (nD) 数组如何在线性 (1D) 内存地址空间中存储的数据表示。内存格式的概念包含两个方面：
- 物理顺序 (Physical Order) 是指数据在物理内存中的存储布局。对于视觉模型，我们通常谈论 NCHW 和 NHWC。这些是对物理内存布局的描述，分别被称为 Channels First（通道在前）和 Channels Last（通道在后）。
- 逻辑顺序 (Logical Order) 是一种描述张量形状和步长的约定。在 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 深度神经网络库) 以在英特尔 CPU 上实现最佳性能。由于直接使用 Channels First 内存格式在物理上无法实现高度优化的性能，因此输入和权重首先会被转换为块格式（blocked format），然后再进行计算。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 的支持范围对于实现最佳性能至关重要。我们已经为计算机视觉领域常用的算子实现了 Channels Last 内核，适用于推理和训练，例如：
- 激活函数（例如 ReLU, PReLU 等）
- 卷积（例如 Conv2d）
- 归一化（例如 BatchNorm2d, GroupNorm 等）
- 池化（例如 AdaptiveAvgPool2d, MaxPool2d 等）
- 重排（例如 ChannelShuffle, PixelShuffle）
详情请参阅 Operators-with-Channels-Last-support。
Channels Last 的原生级优化
如上所述，PyTorch 使用 oneDNN 在英特尔 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 即可。以下是一个使用 TorchVision 在 Channels Last 内存格式下运行 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 一起使用。
性能增益
我们在英特尔® 至强®铂金 8380 CPU @ 2.3 GHz 上对 TorchVision 模型进行了基准测试，每个插槽单实例（batch size = 2 x 物理核心数）。结果显示，Channels Last 比 Channels First 性能提升了 1.3 倍到 1.8 倍。
性能增益主要来自两个方面：
- 对于卷积层，Channels Last 省去了激活函数转换为块格式的步骤，从而提高了整体计算效率。
- 对于池化和上采样层，Channels Last 可以在最内层维度（即“C”）使用向量化逻辑，而 Channels First 则不能。
对于不感知内存格式的层，Channels Last 和 Channels First 具有相同的性能。
结论与未来工作
在本博客中，我们介绍了 Channels Last 的基本概念，并展示了在 CPU 上使用 Channels Last 对视觉模型带来的性能优势。目前的工作阶段仅限于 2D 模型，我们将在不久的将来将优化工作扩展到 3D 模型！
致谢
本博客中介绍的成果是 Meta 和英特尔 PyTorch 团队共同努力的结果。特别感谢来自 Meta 的 Vitaly Fedyunin 和 Wei Wei，他们投入了宝贵的时间并提供了巨大的帮助！我们共同朝着改善 PyTorch CPU 生态系统的道路上又迈出了一步。