quick_test_pytorch.py

import time

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader

if torch.__version__ < '2.4':
    try:
        import intel_extension_for_pytorch as ipex
        print("Intel Extension for PyTorch installed, Version:", ipex.__version__)
    except ImportError:
        print("Failed to import Intel Extension for PyTorch.")

device = torch.device('xpu:0')
device_name = torch.xpu.get_device_name(0)
device_properties = torch.xpu.get_device_properties(0)

print(f"PyTorch Version: {torch.__version__}")
print(f"Using device: {device}")
print(f"Device Name: {device_name}")
print(f"Device Properties: {device_properties}")

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_data = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_data = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64, shuffle=False)
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 10)
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
num_epochs = 5

for epoch in range(num_epochs):
    model.train()
    start_time = time.time()
    running_loss = 0.0
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * images.size(0)
        if i % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')
    epoch_duration = time.time() - start_time
    epoch_loss = running_loss / len(train_loader.dataset)
    print(f'Epoch [{epoch+1}/{num_epochs}] completed in {epoch_duration:.2f} sec, Loss: {epoch_loss:.4f}')

model.eval()
start_time = time.time()
correct = 0
total = 0
for images, labels in test_loader:
    images, labels = images.to(device), labels.to(device)
    outputs = model(images)
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
test_duration = time.time() - start_time
print(f'Evaluated in {test_duration:.2f} sec, Accuracy on test images: {accuracy:.2f}%')