注意
点击 此处 下载完整的示例代码
英特尔® PyTorch* 扩展¶
作者:徐静
英特尔 PyTorch* 扩展通过针对英特尔硬件的优化扩展了 PyTorch,从而进一步提升了性能。大多数优化最终都将包含在 PyTorch 的正式版本中,该扩展的目的是为英特尔硬件上的 PyTorch 提供最新的功能和优化,例如 AVX-512 矢量神经网络指令 (AVX512 VNNI) 和英特尔® 高级矩阵扩展 (英特尔® AMX)。
英特尔® PyTorch* 扩展已作为开源项目在 Github 上发布。
功能¶
- 易于使用的 Python API:英特尔® PyTorch* 扩展为用户提供了简单的前端 Python API 和实用程序,只需少量代码更改即可获得性能优化,例如图优化和算符优化。通常,只需在原始代码中添加 2 到 3 个子句即可。
- 通道最后:与默认的 NCHW 内存格式相比,通道最后 (NHWC) 内存格式可以进一步加速卷积神经网络。在英特尔® PyTorch* 扩展中,NHWC 内存格式已针对大多数关键 CPU 算符启用,但并非所有算符都已合并到 PyTorch 主分支中。预计它们很快就会完全登陆 PyTorch 上游。
- 自动混合精度 (AMP):低精度数据类型 BFloat16 已在第三代至强可扩展服务器(代号 Cooper Lake)上得到原生支持,该服务器采用 AVX512 指令集,并将在下一代英特尔® 至强® 可扩展处理器上得到支持,该处理器采用英特尔® 高级矩阵扩展 (英特尔® AMX) 指令集,性能将进一步提升。英特尔® PyTorch* 扩展已全面支持 CPU 的 BFloat16 自动混合精度 (AMP) 以及算符的 BFloat16 优化,并已部分上游到 PyTorch 主分支。大多数优化将通过正在提交和审查的 PR 登陆 PyTorch 主版本。
- 图优化:为了使用 torchscript 进一步优化性能,英特尔® PyTorch* 扩展支持融合常用算符模式,例如 Conv2D+ReLU、Linear+ReLU 等。融合的好处以透明的方式传递给用户。有关支持的详细融合模式,请参见 此处。图优化将随着 oneDNN Graph API 的引入而上游到 PyTorch。
- 算符优化:英特尔® PyTorch* 扩展还优化了算符,并为性能实现了多个自定义算符。英特尔® PyTorch* 扩展通过 ATen 注册机制用其优化的对应算符替换了一些 ATen 算符。此外,还针对几种常见的拓扑结构实现了一些自定义算符。例如,ROIAlign 和 NMS 在 Mask R-CNN 中定义。为了提高这些拓扑结构的性能,英特尔® PyTorch* 扩展还优化了这些自定义算符。
入门¶
用户只需进行少量代码更改即可开始使用英特尔® PyTorch* 扩展。支持 PyTorch 命令式模式和 TorchScript 模式。本节介绍如何在命令式模式和 TorchScript 模式下使用英特尔® PyTorch* 扩展 API 函数,涵盖数据类型 Float32 和 BFloat16。最后还将介绍 C++ 的用法。
您只需要导入英特尔® PyTorch* 扩展包,并将其优化函数应用于模型对象。如果是训练工作负载,则还需要将优化函数应用于优化器对象。
为了使用 BFloat16 数据类型进行训练和推理,PyTorch 上游已启用 torch.cpu.amp 以方便地支持混合精度,并且 PyTorch 上游和英特尔® PyTorch* 扩展已为 CPU 算符广泛启用 BFloat16 数据类型。运行 torch.cpu.amp 将为每个算符匹配其适当的数据类型,并返回最佳性能。
英特尔® PyTorch* 扩展所需的代码更改在上面一行的注释中突出显示。
训练¶
Float32¶
import torch
import torch.nn as nn
# Import intel_extension_for_pytorch
import intel_extension_for_pytorch as ipex
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = nn.Linear(4, 5)
def forward(self, input):
return self.linear(input)
model = Model()
model.set_state_dict(torch.load(PATH))
optimizer.set_state_dict(torch.load(PATH))
# Invoke optimize function against the model object and optimizer object
model, optimizer = ipex.optimize(model, optimizer, dtype=torch.float32)
for images, label in train_loader():
# Setting memory_format to torch.channels_last could improve performance with 4D input data. This is optional.
images = images.to(memory_format=torch.channels_last)
loss = criterion(model(images), label)
loss.backward()
optimizer.step()
torch.save(model.state_dict(), PATH)
torch.save(optimizer.state_dict(), PATH)
BFloat16¶
import torch
import torch.nn as nn
# Import intel_extension_for_pytorch
import intel_extension_for_pytorch as ipex
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = nn.Linear(4, 5)
def forward(self, input):
return self.linear(input)
model = Model()
model.set_state_dict(torch.load(PATH))
optimizer.set_state_dict(torch.load(PATH))
# Invoke optimize function against the model object and optimizer object with data type set to torch.bfloat16
model, optimizer = ipex.optimize(model, optimizer, dtype=torch.bfloat16)
for images, label in train_loader():
with torch.cpu.amp.autocast():
# Setting memory_format to torch.channels_last could improve performance with 4D input data. This is optional.
images = images.to(memory_format=torch.channels_last)
loss = criterion(model(images), label)
loss.backward()
optimizer.step()
torch.save(model.state_dict(), PATH)
torch.save(optimizer.state_dict(), PATH)
推理 - 命令式模式¶
import torch
import torch.nn as nn
# Import intel_extension_for_pytorch
import intel_extension_for_pytorch as ipex
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = nn.Linear(4, 5)
def forward(self, input):
return self.linear(input)
input = torch.randn(2, 4)
model = Model()
model.eval()
# Invoke optimize function against the model object
model = ipex.optimize(model, dtype=torch.float32)
res = model(input)
BFloat16¶
import torch
import torch.nn as nn
# Import intel_extension_for_pytorch
import intel_extension_for_pytorch as ipex
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = nn.Linear(4, 5)
def forward(self, input):
return self.linear(input)
input = torch.randn(2, 4)
model = Model()
model.eval()
# Invoke optimize function against the model object with data type set to torch.bfloat16
model = ipex.optimize(model, dtype=torch.bfloat16)
with torch.cpu.amp.autocast():
res = model(input)
推理 - TorchScript 模式¶
TorchScript 模式使图优化成为可能,因此提高了一些拓扑结构的性能。英特尔® PyTorch* 扩展支持最常用的算符模式融合,用户无需更改代码即可获得性能提升。
import torch
import torch.nn as nn
# Import intel_extension_for_pytorch
import intel_extension_for_pytorch as ipex
# oneDNN graph fusion is enabled by default, uncomment the line below to disable it explicitly
# ipex.enable_onednn_fusion(False)
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = nn.Linear(4, 5)
def forward(self, input):
return self.linear(input)
input = torch.randn(2, 4)
model = Model()
model.eval()
# Invoke optimize function against the model object
model = ipex.optimize(model, dtype=torch.float32)
model = torch.jit.trace(model, torch.randn(2, 4))
model = torch.jit.freeze(model)
res = model(input)
BFloat16¶
import torch
import torch.nn as nn
# Import intel_extension_for_pytorch
import intel_extension_for_pytorch as ipex
# oneDNN graph fusion is enabled by default, uncomment the line below to disable it explicitly
# ipex.enable_onednn_fusion(False)
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = nn.Linear(4, 5)
def forward(self, input):
return self.linear(input)
input = torch.randn(2, 4)
model = Model()
model.eval()
# Invoke optimize function against the model with data type set to torch.bfloat16
model = ipex.optimize(model, dtype=torch.bfloat16)
with torch.cpu.amp.autocast():
model = torch.jit.trace(model, torch.randn(2, 4))
model = torch.jit.freeze(model)
res = model(input)
C++¶
为了使用 PyTorch 的 C++ 库 libtorch,英特尔® PyTorch* 扩展也提供了其 C++ 动态库。该 C++ 库仅用于处理推理工作负载,例如服务部署。对于常规开发,请使用 Python 接口。与 libtorch 的用法相比,除了将输入数据转换为通道最后数据格式外,无需进行任何特定的代码更改。编译遵循 CMake 的推荐方法。有关详细说明,请参见 PyTorch 教程。在编译过程中,一旦链接了英特尔® PyTorch* 扩展的 C++ 动态库,英特尔优化就会自动激活。
example-app.cpp
'''
#include <torch/script.h>
#include <iostream>
#include <memory>
int main(int argc, const char* argv[]) {
torch::jit::script::Module module;
try {
module = torch::jit::load(argv[1]);
}
catch (const c10::Error& e) {
std::cerr << "error loading the model\n";
return -1;
}
std::vector<torch::jit::IValue> inputs;
// make sure input data are converted to channels last format
inputs.push_back(torch::ones({1, 3, 224, 224}).to(c10::MemoryFormat::ChannelsLast));
at::Tensor output = module.forward(inputs).toTensor();
return 0;
}
'''
CMakeList.txt
'''
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example-app)
find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS} -Wl,--no-as-needed")
add_executable(example-app example-app.cpp)
# Link the binary against the C++ dynamic library file of Intel® Extension for PyTorch*
target_link_libraries(example-app "${TORCH_LIBRARIES}" "${INTEL_EXTENSION_FOR_PYTORCH_PATH}/lib/libintel-ext-pt-cpu.so")
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)
'''
**注意:** 由于英特尔® PyTorch* 扩展仍在开发中,主分支中 C++ 动态库的名称可能与上面显示的 *libintel-ext-pt-cpu.so* 不同。请在安装文件夹中查看名称。so 文件名以 *libintel-* 开头。
编译命令
'''
cmake -DCMAKE_PREFIX_PATH=<LIBPYTORCH_PATH> -DINTEL_EXTENSION_FOR_PYTORCH_PATH=<INTEL_EXTENSION_FOR_PYTORCH_INSTALLATION_PATH> ..
'''