使用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/

推荐阅读

深入理解生成模型VAE

DropBlock的原理和实现

SOTA模型Swin Transformer是如何炼成的!

有码有颜！你要的生成模型VQ-VAE来了！

集成YYDS！让你的模型更快更准！

辅助模块加速收敛，精度大幅提升！移动端实时的NanoDet-Plus来了！

SimMIM：一种更简单的MIM方法

SSD的torchvision版本实现详解

机器学习算法工程师

一个用心的公众号