Intel® Extension for PyTorch* 後端於 Intel® CPU 上

建立於：2023 年 10 月 03 日 | 最後更新：2024 年 06 月 11 日 | 最後驗證：2024 年 11 月 05 日

為了在 Intel® CPU 上更好地使用 torch.compile，Intel® Extension for PyTorch* 實作了一個 ipex 後端。 它的目標是提高 Intel 平台上硬體資源使用的效率，以獲得更好的效能。ipex 後端使用 Intel® Extension for PyTorch* 中設計的進一步客製化來實作模型編譯。

使用範例

訓練 FP32

查看下面的範例，了解如何使用 ipex 後端與 torch.compile 進行具有 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

查看下面的範例，了解如何使用 ipex 後端與 torch.compile 進行具有 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

查看下面的範例，了解如何使用 ipex 後端與 torch.compile 進行具有 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

查看下面的範例，了解如何使用 ipex 後端與 torch.compile 進行具有 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)