Intel® Extension for PyTorch* Backend on Intel® CPUs
Created On: Oct 03, 2023 | Last Updated: Jun 11, 2024 | Last Verified: Nov 05, 2024
To work better with torch.compile on Intel® CPUs, Intel® Extension for PyTorch* implements a backend ipex
.
It targets to improve hardware resource usage efficiency on Intel platforms for better performance.
The ipex backend is implemented with further customizations designed in Intel® Extension for
PyTorch* for the model compilation.
Usage Example
Train FP32
Check the example below to learn how to utilize the ipex backend with torch.compile for model training with FP32 data type.
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()
Train BF16
Check the example below to learn how to utilize the ipex backend with torch.compile for model training with BFloat16 data type.
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()
Inference FP32
Check the example below to learn how to utilize the ipex backend with torch.compile for model inference with FP32 data type.
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)
Inference BF16
Check the example below to learn how to utilize the ipex backend with torch.compile for model inference with BFloat16 data type.
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)