使用Channel Last在CPU上加速PyTorch视觉模型
点蓝色字关注
“机器学习算法工程师
”
设为 星标 ，干货直达！
概述
内存格式在运行视觉模型时对性能有显着影响，通常从性能角度来看，由于更好的数据局部性，Channels Last 更有利。本博客将介绍内存格式的基本概念，并展示使用 Channels Last 在Intel® Xeon® 可扩展处理器上运行流行的 PyTorch 视觉模型上的性能优势。
内存格式介绍
内存格式是指描述多维 (nD) 数组如何存储在线性 (1D) 内存地址空间中的数据表示。内存格式的概念有两个方面：
- 物理顺序：是物理内存中数据存储的布局。对于视觉模型，通常我们谈论的是NCHW、NHWC。这些是物理内存布局的描述，也分别称为 Channels First 和 Channels Last。
- 逻辑顺序：是关于如何描述张量shape和stride的约定。在 PyTorch 中，这个约定是 NCHW。无论物理顺序是什么，张量shape和stride都将始终按照 NCHW 的顺序进行描述。
图 1 是一个shape为 [1, 3, 4, 4] 的张量在 Channels First 和 Channels Last 内存格式上的物理内存布局（通道分别表示为 R、G、B）：
内存格式变换
PyTorch 内存格式传播的一般规则是保留输入张量的内存格式。这意味着 Channels First 输入将生成 Channels First 输出，Channels Last 输入将生成 Channels Last 输出。对于卷积层，PyTorch 默认使用 oneDNN（oneAPI 深度神经网络库）在 Intel CPU 上实现最佳性能。由于使用 Channels Frist 内存格式直接实现高度优化的性能在物理上是不可能的，因此首先将输入和权重转换为块格式，然后进行计算。oneDNN 可以根据输入形状、数据类型和硬件架构选择不同的分块格式，用于矢量化和缓存重用目的。阻塞格式对 PyTorch 是不透明的，因此需要先将输出转换回 Channels。虽然阻塞格式会带来最佳的计算性能，但格式转换可能会增加开销，从而抵消性能增益。另一方面，oneDNN 针对 Channels Last 内存格式进行了优化，以直接使用它以获得最佳性能，PyTorch 将简单地将内存视图传递给 oneDNN。这意味着输入和输出张量的转换被保存。图 2 表示卷积在 PyTorch CPU 上的内存格式变换行为（实线箭头表示内存格式转换，虚线箭头表示内存视图）：
在 PyTorch 上，默认的内存格式是 Channels First。如果特定操作不支持 Channels Last，则 NHWC 输入将被视为非连续 NCHW，因此回退到 Channels First，这将消耗 CPU 上先前的内存带宽并导致性能欠佳。
因此，扩展 Channels Last 支持的范围以获得最佳性能非常重要。我们为 CV 领域的常用算子实现了 Channels Last 内核，适用于推理和训练，例如：
- Activations (e.g., ReLU, PReLU, etc.)
- Convolution (e.g., Conv2d)
- Normalization (e.g., BatchNorm2d, GroupNorm, etc.)
- Pooling (e.g., AdaptiveAvgPool2d, MaxPool2d, etc.)
- Shuffle (e.g., ChannelShuffle, PixelShuffle)
更多详情请参考Operators-with-Channels-Last-support。
Channel Last的原生级别优化
如上所述，PyTorch 使用 oneDNN 在 Intel CPU 上实现卷积的最佳性能。其余的内存格式感知运算符在 PyTorch 本机级别进行了优化，不需要任何第三方库支持。
- 缓存友好的并行化方案：为所有内存格式感知运算符保持相同的并行化方案，这将有助于在将每一层的输出传递到下一层时增加数据局部性。
- 多架构向量化：通常，我们可以在 Channels Last 内存格式的最内部维度上进行向量化。并且将为 AVX2 和 AVX512 生成每个矢量化 CPU 内核。
在为 Channels Last 内核做出贡献的同时，我们也尽最大努力优化 Channels First 对应的内核。事实上，某些算子在物理上不可能在 Channels First 上实现最佳性能，例如卷积、池化等。
使用Channel Last运行视觉模型
在 PyTorch 内存格式教程（https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html）中记录了 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 即可。以下是一个最小示例，展示了如何在 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 脚本）与 Channels Last 一起应用。
性能增益
我们在 Intel® Xeon® Platinum 8380 CPU @ 2.3 GHz、每个插槽单个实例（批量大小 = 2 x 物理内核数）上对 TorchVision 模型的推理性能进行了基准测试。结果显示，Channels Last 的性能比 Channels First 高 1.3 到 1.8 倍。
性能提升主要来自两个方面：
- 对于卷积层，Channels Last 保存了内存格式转换为块格式的激活，从而提高了整体计算效率。
- 对于 Pooling 和 Upsampling 层，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 生态系统的道路上又迈出了一步。
本文翻译自https://pytorch.org/blog/accelerating-pytorch-vision-models-with-channels-last-on-cpu/
推荐阅读
辅助模块加速收敛，精度大幅提升！移动端实时的NanoDet-Plus来了！
机器学习算法工程师
一个用心的公众号
