评价此页

★ ★ ★ ★ ★

recipes/intel_extension_for_pytorch

在 Google Colab 中运行

Intel® PyTorch 扩展*#

创建日期：2021年11月09日 | 最后更新：2025年07月24日 | 最后验证：2024年11月05日

Intel® PyTorch 扩展* 为 PyTorch* 增加了最新的功能和优化，可在 Intel 硬件上获得额外的性能提升。这些优化利用了 Intel CPU 上的 AVX-512 向量神经网络指令 (AVX512 VNNI) 和 Intel® 高级矩阵扩展 (Intel® AMX)，以及 Intel 独立 GPU 上的 Intel X^e 矩阵扩展 (XMX) AI 引擎。此外，通过 PyTorch* 的 xpu 设备，Intel® PyTorch 扩展* 为具有 PyTorch* 的 Intel 独立 GPU 提供了简便的 GPU 加速。

Intel® PyTorch 扩展* 已作为开源项目发布在 Github 上。

CPU 的源代码可在 main 分支上找到。
GPU 的源代码可在 xpu-main 分支上找到。

特性#

Intel® PyTorch 扩展* 在 CPU 和 GPU 上共享大多数特性。

易于使用的 Python API： Intel® PyTorch 扩展* 提供简单的前端 Python API 和实用工具，供用户通过少量代码修改即可获得图优化和算子优化等性能优化。通常，只需要在原始代码中添加 2 到 3 条语句。
通道优先（Channels Last）： 与默认的 NCHW 内存格式相比，通道优先 (NHWC) 内存格式可以进一步加速卷积神经网络。在 Intel® PyTorch 扩展* 中，NHWC 内存格式已为大多数关键 CPU 算子启用，尽管并非所有算子都已合并到 PyTorch 主分支。预计很快将完全合并到 PyTorch 主线。
自动混合精度 (AMP)： 低精度数据类型 BFloat16 在 3 代 Xeon 可扩展服务器 (又名 Cooper Lake) 上通过 AVX512 指令集得到原生支持，并且将在下一代 Intel® Xeon® 可扩展处理器上通过 Intel® 高级矩阵扩展 (Intel® AMX) 指令集得到支持，从而提供更强的性能。Intel® PyTorch 扩展* 中已大规模启用 CPU 的 BFloat16 自动混合精度 (AMP) 以及 BFloat16 算子优化，并部分合并到 PyTorch 主分支。其中大多数优化将通过正在提交和审查的 PR 合并到 PyTorch 主分支。Intel 独立 GPU 已启用 BFloat16 和 Float16 的自动混合精度 (AMP)。
图优化： 为了进一步优化性能，Intel® PyTorch 扩展* 支持对常用算子模式进行融合，例如 Conv2D+ReLU、Linear+ReLU 等。融合的好处将以透明的方式提供给用户。支持的详细融合模式可以在此处找到。图优化将通过 oneDNN 图 API 的引入合并到 PyTorch 中。
算子优化： Intel® PyTorch 扩展* 还优化算子并实现了一些定制算子以提高性能。一些 ATen 算子通过 ATen 注册机制替换为 Intel® PyTorch 扩展* 中经过优化的对应算子。此外，为几个流行的拓扑结构实现了一些定制算子。例如，ROIAlign 和 NMS 在 Mask R-CNN 中定义。为了提高这些拓扑结构的性能，Intel® PyTorch 扩展* 还优化了这些定制算子。

入门#

用户只需进行少量代码修改即可开始使用 Intel® PyTorch 扩展*。支持 PyTorch 命令式模式和 TorchScript 模式。本节介绍 Intel® PyTorch 扩展* API 函数在命令式模式和 TorchScript 模式下的用法，涵盖 Float32 和 BFloat16 数据类型。最后还将介绍 C++ 的用法。

您只需要导入 Intel® PyTorch 扩展* 包并将其优化函数应用于模型对象。如果是训练工作负载，还需要将优化函数应用于优化器对象。

对于使用 BFloat16 数据类型的训练和推理，PyTorch 主线已启用 torch.cpu.amp 以方便支持混合精度。PyTorch 主线和 Intel® PyTorch 扩展* 已广泛启用 CPU 算子的 BFloat16 数据类型。同时，由 Intel® PyTorch 扩展* 注册的 torch.xpu.amp 使得在 Intel 独立 GPU 上可以方便地使用 BFloat16 和 Float16 数据类型。 torch.cpu.amp 或 torch.xpu.amp 会自动将每个算子匹配到其相应的数据类型，并返回最佳性能。

示例 – CPU#

本节展示在 CPU 上使用 Intel® PyTorch 扩展* 进行训练和推理的示例

Intel® PyTorch 扩展* 所需的代码更改已突出显示。

训练#

Float32#

import torch
import torchvision
import intel_extension_for_pytorch as ipex

LR = 0.001
DOWNLOAD = True
DATA = 'datasets/cifar10/'

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = torchvision.datasets.CIFAR10(
        root=DATA,
        train=True,
        transform=transform,
        download=DOWNLOAD,
)
train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=128
)

model = torchvision.models.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr = LR, momentum=0.9)
model.train()
model, optimizer = ipex.optimize(model, optimizer=optimizer)

for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print(batch_idx)
torch.save({
     'model_state_dict': model.state_dict(),
     'optimizer_state_dict': optimizer.state_dict(),
     }, 'checkpoint.pth')

BFloat16#

import torch
import torchvision
import intel_extension_for_pytorch as ipex

LR = 0.001
DOWNLOAD = True
DATA = 'datasets/cifar10/'

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = torchvision.datasets.CIFAR10(
        root=DATA,
        train=True,
        transform=transform,
        download=DOWNLOAD,
)
train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=128
)

model = torchvision.models.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr = LR, momentum=0.9)
model.train()
model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.bfloat16)

for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    with torch.cpu.amp.autocast():
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
    optimizer.step()
    print(batch_idx)
torch.save({
     'model_state_dict': model.state_dict(),
     'optimizer_state_dict': optimizer.state_dict(),
     }, 'checkpoint.pth')

推理#

Float32#

import torch
import torchvision.models as models

model = models.resnet50(pretrained=True)
model.eval()
data = torch.rand(1, 3, 224, 224)

#################### code changes ####################
import intel_extension_for_pytorch as ipex
model = ipex.optimize(model)
######################################################

with torch.no_grad():
  model(data)

BFloat16#

import torch
from transformers import BertModel

model = BertModel.from_pretrained(args.model_name)
model.eval()

vocab_size = model.config.vocab_size
batch_size = 1
seq_length = 512
data = torch.randint(vocab_size, size=[batch_size, seq_length])

#################### code changes ####################
import intel_extension_for_pytorch as ipex
model = ipex.optimize(model, dtype=torch.bfloat16)
######################################################

with torch.no_grad():
  with torch.cpu.amp.autocast():
    model(data)

示例 – GPU#

本节展示在 GPU 上使用 Intel® PyTorch 扩展* 进行训练和推理的示例

所需的 Intel® PyTorch 扩展* 代码更改已通过注释行突出显示。

训练#

Float32#

import torch
import torchvision
############# code changes ###############
import intel_extension_for_pytorch as ipex
############# code changes ###############

LR = 0.001
DOWNLOAD = True
DATA = 'datasets/cifar10/'

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = torchvision.datasets.CIFAR10(
        root=DATA,
        train=True,
        transform=transform,
        download=DOWNLOAD,
)
train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=128
)

model = torchvision.models.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr = LR, momentum=0.9)
model.train()
#################################### code changes ################################
model = model.to("xpu")
model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.float32)
#################################### code changes ################################

for batch_idx, (data, target) in enumerate(train_loader):
    ########## code changes ##########
    data = data.to("xpu")
    target = target.to("xpu")
    ########## code changes ##########
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print(batch_idx)
torch.save({
     'model_state_dict': model.state_dict(),
     'optimizer_state_dict': optimizer.state_dict(),
     }, 'checkpoint.pth')

BFloat16#

import torch
import torchvision
############# code changes ###############
import intel_extension_for_pytorch as ipex
############# code changes ###############

LR = 0.001
DOWNLOAD = True
DATA = 'datasets/cifar10/'

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = torchvision.datasets.CIFAR10(
        root=DATA,
        train=True,
        transform=transform,
        download=DOWNLOAD,
)
train_loader = torch.utils.data.DataLoader(
        dataset=train_dataset,
        batch_size=128
)

model = torchvision.models.resnet50()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr = LR, momentum=0.9)
model.train()
##################################### code changes ################################
model = model.to("xpu")
model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.bfloat16)
##################################### code changes ################################

for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    ######################### code changes #########################
    data = data.to("xpu")
    target = target.to("xpu")
    with torch.xpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    ######################### code changes #########################
        output = model(data)
        loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print(batch_idx)
torch.save({
     'model_state_dict': model.state_dict(),
     'optimizer_state_dict': optimizer.state_dict(),
     }, 'checkpoint.pth')

推理#

Float32#

import torch
import torchvision.models as models
############# code changes ###############
import intel_extension_for_pytorch as ipex
############# code changes ###############

model = models.resnet50(pretrained=True)
model.eval()
data = torch.rand(1, 3, 224, 224)

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

#################### code changes ################
model = model.to("xpu")
data = data.to("xpu")
model = ipex.optimize(model, dtype=torch.float32)
#################### code changes ################

with torch.no_grad():
  model(data)

BFloat16#

import torch
import torchvision.models as models
############# code changes ###############
import intel_extension_for_pytorch as ipex
############# code changes ###############

model = models.resnet50(pretrained=True)
model.eval()
data = torch.rand(1, 3, 224, 224)

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

#################### code changes #################
model = model.to("xpu")
data = data.to("xpu")
model = ipex.optimize(model, dtype=torch.bfloat16)
#################### code changes #################

with torch.no_grad():
  ################################# code changes ######################################
  with torch.xpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=False):
  ################################# code changes ######################################
    model(data)

Float16#

import torch
import torchvision.models as models
############# code changes ###############
import intel_extension_for_pytorch as ipex
############# code changes ###############

model = models.resnet50(pretrained=True)
model.eval()
data = torch.rand(1, 3, 224, 224)

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

#################### code changes ################
model = model.to("xpu")
data = data.to("xpu")
model = ipex.optimize(model, dtype=torch.float16)
#################### code changes ################

with torch.no_grad():
  ################################# code changes ######################################
  with torch.xpu.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=False):
  ################################# code changes ######################################
    model(data)

C++ (仅限 CPU)#

为了与 PyTorch 的 C++ 库 libtorch 一起工作，Intel® PyTorch 扩展* 也提供了其 C++ 动态库。C++ 库主要用于处理推理工作负载，例如服务部署。对于常规开发，请使用 Python 接口。与使用 libtorch 相比，不需要进行特定的代码更改，只需将输入数据转换为通道优先的数据格式。编译遵循 CMake 的推荐方法。详细说明可以在 PyTorch 教程中找到。在编译过程中，一旦链接了 Intel® PyTorch 扩展* 的 C++ 动态库，Intel 优化将自动激活。

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::rand({1, 3, 224, 224});

    at::Tensor output = module.forward(inputs).toTensor();
    std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << std::endl;
    std::cout << "Execution finished" << std::endl;
    return 0;
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example-app)

find_package(intel_ext_pt_cpu REQUIRED)

add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")

set_property(TARGET example-app PROPERTY CXX_STANDARD 17)

编译命令

$ cmake -DCMAKE_PREFIX_PATH=<LIBPYTORCH_PATH> ..
$ make

如果显示 Found INTEL_EXT_PT_CPU 为 TRUE，则表示该扩展已链接到二进制文件中。可以使用 Linux 命令 ldd 进行验证。

$ cmake -DCMAKE_PREFIX_PATH=/workspace/libtorch ..
-- The C compiler identification is GNU XX.X.X
-- The CXX compiler identification is GNU XX.X.X
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Check for working C compiler: /usr/bin/cc - skipped
-- Detecting C compile features
-- Detecting C compile features - done
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Check for working CXX compiler: /usr/bin/c++ - skipped
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Found Torch: /workspace/libtorch/lib/libtorch.so
-- Found IPEX: /workspace/libtorch/lib/libintel-ext-pt-cpu.so
-- Configuring done
-- Generating done
-- Build files have been written to: /workspace/build

$ ldd example-app
        ...
        libtorch.so => /workspace/libtorch/lib/libtorch.so (0x00007f3cf98e0000)
        libc10.so => /workspace/libtorch/lib/libc10.so (0x00007f3cf985a000)
        libintel-ext-pt-cpu.so => /workspace/libtorch/lib/libintel-ext-pt-cpu.so (0x00007f3cf70fc000)
        libtorch_cpu.so => /workspace/libtorch/lib/libtorch_cpu.so (0x00007f3ce16ac000)

教程#

更详细的教程可在官方 Intel® PyTorch 扩展* 文档中找到

CPU
GPU

Intel® PyTorch 扩展*#

特性#

入门#

示例 – CPU#

训练#

Float32#

BFloat16#

推理#

Float32#

BFloat16#

示例 – GPU#

训练#

Float32#

BFloat16#

推理#

Float32#

BFloat16#

Float16#

C++ (仅限 CPU)#

教程#

文档

教程

资源