快捷方式

英特尔® 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> ..
'''

文档

获取 PyTorch 全面的开发者文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源