概述
内存格式在运行视觉模型时对性能有显著影响，通常从性能角度来看，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 内存格式在物理上无法直接实现高度优化的性能，输入和权重首先会被转换为 blocked format 然后进行计算。oneDNN 可能会根据输入形状、数据类型和硬件架构选择不同的 blocked format，以实现向量化和缓存重用。Blocked format 对于 PyTorch 是不透明的，因此输出需要转换回 Channels First。尽管 blocked format 可以带来最佳计算性能，但格式转换可能会增加开销，并因此抵消性能提升。
另一方面，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)
参考 支持 Channels Last 的算子 了解详情。
Channels Last 的原生级别优化
如上所述，PyTorch 使用 oneDNN 在 Intel CPU 上为卷积实现最佳性能。其余支持内存格式的算子在 PyTorch 原生级别进行了优化，这不依赖任何第三方库的支持。
- 缓存友好的并行化方案： 对所有支持内存格式的算子保持相同的并行化方案，这将有助于在将每一层的输出传递给下一层时提高数据局部性。
- 跨多种架构的向量化： 通常，在 Channels Last 内存格式上，我们可以沿着最内层维度（例如“C”）进行向量化。并且每个向量化的 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 一起使用。
性能提升
我们对 TorchVision 模型在 Intel® Xeon® Platinum 8380 CPU @ 2.3 GHz 上的推理性能进行了基准测试，每个 socket 单独一个实例（批处理大小 = 2 x 物理核心数）。结果表明，相较于 Channels First，Channels Last 具有 1.3 倍至 1.8 倍的性能提升。
性能提升主要来自两个方面：
- 对于卷积层，Channels Last 避免了将激活转换为 blocked format 的内存格式转换，这提高了整体计算效率。
- 对于池化层和上采样层，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 生态系统的道路上迈出了又一步。
