Intel® 扩展适用于 Intel® CPU 的 PyTorch* 后端

Created On: Oct 03, 2023 | Last Updated: Jun 11, 2024 | Last Verified: Nov 05, 2024

为了在 Intel® CPU 上更好地使用 torch.compile，Intel® Extension for PyTorch* 实现了一个名为 ipex 的后端。它旨在提高 Intel 平台上的硬件资源利用效率，从而实现更好的性能。ipex 后端通过 Intel® Extension for PyTorch* 中设计的进一步自定义功能来实现模型编译。

使用示例

训练 FP32

查看下面的示例，了解如何使用 torch.compile 和 ipex 后端进行 FP32 数据类型的模型训练。

import torch
import torchvision

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 ####################
import intel_extension_for_pytorch as ipex

# Invoke the following API optionally, to apply frontend optimizations
model, optimizer = ipex.optimize(model, optimizer=optimizer)

compile_model = torch.compile(model, backend="ipex")
######################################################

for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    output = compile_model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

训练 BF16

查看下面的示例，了解如何使用 torch.compile 和 ipex 后端进行 BFloat16 数据类型的模型训练。

import torch
import torchvision

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 ####################
import intel_extension_for_pytorch as ipex

# Invoke the following API optionally, to apply frontend optimizations
model, optimizer = ipex.optimize(model, dtype=torch.bfloat16, optimizer=optimizer)

compile_model = torch.compile(model, backend="ipex")
######################################################

with torch.cpu.amp.autocast():
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = compile_model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

推理 FP32

查看下面的示例，了解如何使用 torch.compile 和 ipex 后端进行 FP32 数据类型的模型推理。

import torch
import torchvision.models as models

model = models.resnet50(weights='ResNet50_Weights.DEFAULT')
model.eval()
data = torch.rand(1, 3, 224, 224)

#################### code changes ####################
import intel_extension_for_pytorch as ipex

# Invoke the following API optionally, to apply frontend optimizations
model = ipex.optimize(model, weights_prepack=False)

compile_model = torch.compile(model, backend="ipex")
######################################################

with torch.no_grad():
    compile_model(data)

推理 BF16

查看下面的示例，了解如何使用 torch.compile 和 ipex 后端进行 BFloat16 数据类型的模型推理。

import torch
import torchvision.models as models

model = models.resnet50(weights='ResNet50_Weights.DEFAULT')
model.eval()
data = torch.rand(1, 3, 224, 224)

#################### code changes ####################
import intel_extension_for_pytorch as ipex

# Invoke the following API optionally, to apply frontend optimizations
model = ipex.optimize(model, dtype=torch.bfloat16, weights_prepack=False)

compile_model = torch.compile(model, backend="ipex")
######################################################

with torch.no_grad(), torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    compile_model(data)