概述
内存格式对运行视觉模型时的性能有显著影响，通常从性能角度来看，通道最后（Channels Last）格式更为有利，因为它具有更好的数据局部性。
本文将介绍内存格式的基本概念，并通过在英特尔® 至强® 可扩展处理器上使用通道最后（Channels Last）格式对流行的 PyTorch 视觉模型进行性能提升的演示。
内存格式简介
内存格式指的是描述多维（nD）数组如何在线性（1D）内存地址空间中存储的数据表示。内存格式的概念有两个方面：
- 物理顺序是指数据存储在物理内存中的布局。对于视觉模型，我们通常谈论 NCHW、NHWC。这些都是物理内存布局的描述，也分别被称为通道优先和通道后置。
- 逻辑顺序是描述张量形状和步长的惯例。在 PyTorch 中，这个惯例是 NCHW。无论物理顺序如何，张量形状和步长都将始终以 NCHW 的顺序表示。
图-1 展示了形状为[1, 3, 4, 4]的张量在通道优先和通道后置内存格式（通道分别表示为 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）
请参阅“支持 Channels-Last 的算子”以获取详细信息。
在 Channels Last 上的本地级别优化
如上所述，PyTorch 使用 oneDNN 在 Intel CPU 上实现卷积的最佳性能。其余的内存格式感知算子都在 PyTorch 本地级别进行了优化，无需任何第三方库支持。
- 缓存友好并行化方案：对所有内存格式感知算子保持相同的并行化方案，这有助于在将每一层的输出传递给下一层时提高数据局部性。
- 在多个架构上实现向量化：通常，我们可以在 Channels Last 内存格式上对最内层维度进行向量化。每个向量化 CPU 内核都将为 AVX2 和 AVX512 生成。
在为 Channels Last 内核做出贡献的同时，我们尽力优化 Channels First 对应的版本。事实上，某些算子在 Channels First 上实现最佳性能在物理上是不可行的，例如卷积、池化等。
在 Channels Last 上运行视觉模型
PyTorch 内存格式教程中记录了 Channels Last 相关的 API。通常，我们可以通过以下方式将 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 运行 ResNet50 的最小示例，展示了如何在 Channels Last 内存格式上运行：
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 倍）。结果显示，Channels Last 比 Channels First 有 1.3 倍到 1.8 倍的性能提升。
性能提升主要来自两个方面：
- 对于卷积层，Channels Last 节省了激活函数到阻塞格式的内存格式转换，从而提高了整体计算效率。
- 对于池化和上采样层，Channels Last 可以在最内维度的向量逻辑中使用，例如“C”，而 Channels First 则不能。
对于不关心内存格式的层，通道最后和通道第一的性能相同。
结论与未来工作
在这篇博客中，我们介绍了通道最后的基本概念，并展示了在视觉模型上使用通道最后对 CPU 的性能提升。目前的工作仅限于当前阶段的 2D 模型，我们将在不久的将来将优化工作扩展到 3D 模型！
致谢
本博客中展示的结果是 Meta 和 Intel PyTorch 团队共同努力的结果。特别感谢 Meta 的 Vitaly Fedyunin 和 Wei Wei，他们花费宝贵的时间并给予了实质性的帮助！我们一起在改善 PyTorch CPU 生态系统之路上迈出了新的一步！
