面向 PyTorch* 的英特尔 扩展助力加速 PyTorch

英特尔工程师一直在PyTorch 开源社区积极贡献，以加快PyTorch 在英特尔 CPU 上的运行速度。面向PyTorch* 的英特尔 扩展是英特尔发起的一个开源扩展项目，它基于PyTorch的扩展机制实现，通过提供额外的软件优化极致地发挥硬件特性，帮助用户在原生PyTorch的基础上更最大限度地提升英特尔 CPU 上的深度学习推理计算和训练性能。这些软件优化大部分将会进入未来的原生PyTorch 版本中，但是通过扩展，PyTorch 用户将能更加及时地受益于英特尔硬件的最新功能，并在第一时间体验软件优化带来的最佳性能和部署便利。值得一提的是，除了 CPU，面向 PyTorch* 的英特尔 扩展还将在不久的将来为英特尔 GPU 带来支持和优化。

                                          面向 PyTorch* 的英特尔 扩展概述

面向 PyTorch* 的英特尔 扩展针对命令模式（imperative mode）和图模式进行了优化，并针对 PyTorch 的三个关键模块进行了优化：运算符、图和运行时。优化的运算符和内核通过 PyTorch 调度机制完成调用。在执行期间，面向 PyTorch* 的英特尔 扩展将用扩展中优化的运算符覆盖 ATen 运算符的对应部分，并为常见的用例提供一组额外的自定义运算符和优化实现。在图模式下，扩展进一步应用图优化，以最大限度地提升内核的执行性能。运行时优化封装在运行时扩展模块中，该模块可为用户提供几个 PyTorch 前端 API，以便对线程运行时进行更精细化的控制。 

优化一览

内存布局

内存布局是优化视觉相关运算符的基础，而且使用适合的内存格式来处理输入张量可以显著地提高 PyTorch CNN 模型的性能。有几篇报道（1，2）显示，Channels Last 内存格式通常有利于多个硬件后端，包括英特尔 CPU。 

我们正在针对许多 CNN 相关运算符 实现Channels Last 内核并提交到 PyTorch 上游中，面向 PyTorch* 的英特尔 扩展将提供更全面的支持，而且我们建议用户尽可能使用Channels Last内存格式，详情可参照PyTorch 内存格式教程。

----------------------------------------------------------

model = model.to(memory_format=torch.channels_last)

input = input.to(memory_format=torch.channels_last)

------------------------------------------------------------------------------

英特尔 oneDNN 库针对模型权重引入了块状内存布局，以实现更出色的向量化和缓存复用。为避免运行时转换，在执行 oneDNN 运算符之前，我们将权重转换为预定义的最佳块格式。这种技术被称为权重预封装，当用户调用扩展提供的 ipex.optimize 前端 API 时，它可以用于模型推理和训练。

自定义运算符

在面向 PyTorch* 的英特尔 扩展中，我们提供多个定制的运算符来加速处理常见的拓扑结构，例如Fused Interaction、Merged Embedding Bag等等，用于 DLRM推荐模型和对象检测模型中的ROIAlign、FrozenBatchNorm 等。

优化器（Optimizer）优化

优化器的性能在训练性能中发挥着不可忽视的作用，我们在面向 PyTorch* 的英特尔 扩展中提供优化器的优化版本。

优化器融合

我们支持面向优化器 Lamb、Adagrad 和 SGD 的融合内核，这些融合优化器将用于通过 ipex.optimize 前端优化来替换相应的优化，因此用户无需更改模型代码即可应用这些融合优化器。内核在权重更新步骤中融合模型参数的内存受限操作链及其梯度，以便数据驻留在高速缓存中，无需再从内存中加载。我们正致力于在即将发布的扩展版本中提供更多的融合优化器

优化器中的计算拆分BF16 混合精度训练可通过加速计算、降低内存带宽压力和减少内存消耗显著提升性能然而，随着训练后期不断积累，权重更新会太小而导致精度问题，常见的做法是在 FP32 中保留权重的主副本，这样 BF16 模型权重会增大 2 倍。内存占用率升高会给需要大量权重的工作负载（如推荐模型）造成负担。

我们在应用了一种“拆分”优化，以支持BF16训练。这项技术只需权重增大 1 倍（而非2倍），因此权重大小与进行 FP32 训练时一样。我们将 FP32 参数拆分成“Top half”和“Bottom half”。“Top half”是前 16 位，可被完全视作 BF16 数字，“Bottom half”是后 16 位，其作用是避免出现精度损失。在执行向前和向后传播时，“Top half”仍将受益于英特尔 CPU 上的原生 BF16 支持，而在执行参数更新时，我们将“Top half”和“Bottom half”串级起来，从而将参数恢复成 FP32，以避免精度损失。

低精度数据类型支持

深度学习从业者已证实，低数值精度对于训练和推理计算非常有效，使用 16 位乘法器和 32 位累加器进行训练和推理，精度损失微乎其微或者没有损失，甚至使用 8 位乘法器和 32 位累加器运行推理工作负载，有些工作负载的精度的损失也是完全可接受的，因此推动了它们在工作负载中的采用。

较低精度可通过两种方式增强性能：1) 额外的乘法累加 (MAC) 吞吐量可增强计算受限的操作，2) 使用 16 位（而非 32 位）减少占用空间，可通过减少内存层级结构中的内存事务来增强内存带宽受限的操作。

BF16 自动混合精度

英特尔将原生 BF16 支持引入第三代英特尔 至强 可扩展处理器，并支持 BF16→ FP32 FMA 和 FP32→BF16 转换，以及英特尔 高级矢量扩展指令集（英特尔 AVX-512），相比 FP32 FMA 相比，理论计算吞吐量翻了一番。下一代英特尔 至强 可扩展处理器上的英特尔 高级矩阵扩展（英特尔 AMX）指令集扩展将进一步加速 BF16。

我们已经在原生的 PyTorch发布当中通过API torch.cpu.amp.autocast支持 了CPU 上面向 BF16 的自动混合精度 (AMP)，且运算符的 BF16 优化也已部分登录 包含在原生PyTorch的发布版本当中了。在所有优化进入 PyTorch 上游之前，通过PyTorch* 的英特尔 扩展我们已经实现了大规模的 BF16 运算符支持，用户可提前体验 BF16 的全部优势。

INT8 量化

量化是指通过降低权重和/或激活的数值精度来压缩深度网络中的信息。将用于参数信息的 FP32 数字转换为 INT8 之后，模型会变小，并显著节省内存和降低计算需求。英特尔自第二代英特尔至强可扩展处理器（代号“Cascade Lake”）起开始引入 AVX512 VNNI 扩展指令集，从而支持更快地计算 INT8 数据并提高吞吐量。

PyTorch 提供了几种不同的方法来量化模型（如本文所述）。为了进一步改善用户体验，面向 PyTorch* 的英特尔 扩展提供了一个自动化流程和经过充分验证的默认量化方案。用户通过几个简单的前端 API 就能更轻松地获得 INT8 的性能优势，同时将精度损失降至最低。目前它主要是作为支持训练后静态量化的实验性功能来提供，训练后动态量化支持也即将添加至后续量化过程中。

在量化过程中，将默认执行各种图形优化，比如运算符融合，并且相比原生PyTorch，将会应用更多的 oneDNN 融合功能。它们将逐渐被提交到 stock PyTorch，当前，面向 PyTorch* 的英特尔 扩展将为用户显著提升性能。

图优化

像运算符融合这样的图优化将通过优化整体计算和内存带宽，最大限度地提升基础内核实施的性能。

面向PyTorch* 的英特尔 扩展应用基于TorchScript IR 的算子融合通道，在 oneDNN 提供的算子融合能力和扩展中的融合内核的支持下，用户将能透明地获得算子融合带来的性能提升。

常数折叠是在 PyTorch 前端执行的图优化。它用预先计算的常数节点替换掉部分具备全部常量输入的运算符。例如，用于推理的 Convolution+BatchNorm 折叠，它可将为许多 CNN 模型带来不可忽视的性能优势。

需要注意的一点是，Intel和Meta的工程师正合作致力于PyTorch NNC（神经网络编译器）的开发和优化，通过组合多种算子融合技术，以获得CPU上最佳的图优化性能.

运行时扩展

面向 PyTorch* 的英特尔 扩展中的模块旨在提供运行时优化，例如操作调度、计算资源管理、内存管理。它目前为用户提供了几个 PyTorch 前端 API，以便对线程运行时进行更精细化的控制，这将为多个用例提供优势，比如支持共享权重的多流推理，后续版本将具备更多功能。

程序加载器

正如性能调优指南 (1，2)中所述，为了在 CPU 上实现最佳推理性能，需要尝试和调整一些设置，例如 openMP 库、内存分配器等。在面向 PyTorch* 的英特尔 扩展中，我们提供一个带有默认设置的启动程序脚本，以简化部署，我们建议用户始终使用启动程序脚本启动工作负载。

示例

面向 PyTorch* 的英特尔 扩展可作为面向 Python 程序的 Python 模块加载，也可作为面向 C++ 程序的 C++ 库。用户只需更改几行代码，就可获取该扩展的所有优势。以下列举几个例子，如欲参考更多示例，请参阅我们的教程。

BF16 训练

-----------------------------------------------------------------

...
import torch
import intel_extension_for_pytorch as ipex
...

model = Model()
model = model.to(memory_format=torch.channels_last)

criterion = ...
optimizer = ...
model.train()

model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.bfloat16)
...

with torch.cpu.amp.autocast():
# Setting memory_format to torch.channels_last could improve performance with 4D input data.
data = data.to(memory_format=torch.channels_last)
    optimizer.zero_grad()
    output = model(data)
    loss = ...
    loss.backward()

...

-------------------------------------------------------------------------------

BF16 推理

-----------------------------------------------------------------

...
import torch
import intel_extension_for_pytorch as ipex
...

model = Model()
model = model.to(memory_format=torch.channels_last)

model.eval()

model = ipex.optimize(model, dtype=torch.bfloat16)
...

with torch.cpu.amp.autocast(),torch.no_grad():
    # Setting memory_format to torch.channels_last could improve performance with 4D input data.
    data = data.to(memory_format=torch.channels_last)
    model = torch.jit.trace(model, data)
    model = torch.jit.freeze(model)

with torch.no_grad():
    output = model(data)

...

-------------------------------------------------------------------------------

INT8 推理

标定

---------------------------------------------------------------------------------

import os

import torch

model = Model()
model.eval()
data = torch.rand(<shape>)

# Applying torch.fx.experimental.optimization.fuse against model performs conv-batchnorm folding for better performance.
import torch.fx.experimental.optimization as optimization
model = optimization.fuse(model, inplace=True)

import intel_extension_for_pytorch as ipex
conf = ipex.quantization.QuantConf(qscheme=torch.per_tensor_affine)

for d in calibration_data_loader():
    # conf will be updated with observed statistics during calibrating with the dataset
    with ipex.quantization.calibrate(conf):
        model(d)

conf.save('int8_conf.json', default_recipe=True)
with torch.no_grad():
    model = ipex.quantization.convert(model, conf, torch.rand(<shape>))

model.save('quantization_model.pt')

-------------------------------------------------------------------------------------

部署

-----------------------------------------------------------------

import torch
import intel_extension_for_pytorch as ipex

model = torch.jit.load('quantization_model.pt')
model.eval()
data = torch.rand(<shape>)

with torch.no_grad():
    model(data)

-------------------------------------------------------------------

性能提升

以下列举了通过面向 PyTorch* 英特尔 扩展所获得的部分性能提升。这些数字是在英特尔 至强 铂金 8380 CPU @ 2.3 GHz 上测量所得。离线推理是指使用单路处理器上的所有内核运行大批量单实例推理。实时推理是指运行多实例单批次推理，每个实例有 4 个内核。如欲获取更多详细的测量信息以及 BF16 和 INT8 推断结果，请点击此处。

https://intel.github.io/intel-extension-for-pytorch/1.11/tutorials/performance.html

总结与未来工作

面向 PyTorch* 英特尔 扩展旨在以敏捷、快速的方式帮助用户提升英特尔 CPU 的性能。我们将逐渐地将大部分优化提交至 stock PyTorch，我们还将持续为最新英特尔硬件提供全新的优化和功能，以便用户快速获取和使用。我们强烈建议用户尝试面向 PyTorch* 英特尔 扩展 (github，pypi)，并通过 github向我们发送反馈。同时它也是一个开源项目，我们欢迎社区用户直接通过 github提交pull request。

参考资料

PYTORCH 中的 (BETA) CHANNELS LAST 内存格式

https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html%23sphx-glr-intermediate-memory-format-tutorial-py

· 高效的 PyTorch：张量内存格式至关重要

https://pytorch.org/blog/tensor-memory-format-matters/

· 了解内存格式

https://oneapi-src.github.io/oneDNN/dev_guide_understanding_memory_formats.html

·  PyTorch 中的实际量化

https://pytorch.org/blog/quantization-in-practice/

·   面向 PyTorch 的英特尔扩展文档

https://intel.github.io/intel-extension-for-pytorch/1.11/