Ускорение машинного обучения с помощью графических процессоров в PyTorch

Расширенные примеры использования CUDA в PyTorch

Пример 1: Сравнение скорости CPU и GPU на умножении матриц

Демонстрирует значительное преимущество GPU при работе с большими матрицами.

import torch
import time

size = 3000
A_cpu = torch.randn(size, size)
B_cpu = torch.randn(size, size)

# Замер на CPU
start = time.time()
C_cpu = torch.matmul(A_cpu, B_cpu)
time_cpu = time.time() - start
print(f'CPU: {time_cpu:.3f} сек')

if torch.cuda.is_available():
    A_gpu = A_cpu.cuda()
    B_gpu = B_cpu.cuda()
    torch.cuda.synchronize()  # дождаться завершения
    start = time.time()
    C_gpu = torch.matmul(A_gpu, B_gpu)
    torch.cuda.synchronize()
    time_gpu = time.time() - start
    print(f'GPU: {time_gpu:.3f} сек')
    print(f'Ускорение: {time_cpu/time_gpu:.1f}x')

CPU: 2.145 сек
GPU: 0.025 сек
Ускорение: 85.8x

Пример 2: Обучение сверточной сети на GPU с замером времени

Используется простой классификатор MNIST.

import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import time

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

transform = transforms.ToTensor()
train_dataset = dsets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

model = nn.Sequential(
    nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    nn.Linear(16*14*14, 10)
).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

start = time.time()
for epoch in range(3):
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch+1} завершена')
total = time.time() - start
print(f'Общее время обучения: {total:.2f} сек')

Epoch 1 завершена
Epoch 2 завершена
Epoch 3 завершена
Общее время обучения: 14.23 сек

Пример 3: Использование автоматической смешанной точности (AMP) с обучением

Показывает, как применить AMP для уменьшения расхода памяти и ускорения.

scaler = torch.cuda.amp.GradScaler()

for epoch in range(5):
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()

        with torch.cuda.amp.autocast():
            outputs = model(images)
            loss = criterion(outputs, labels)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

(Вывод аналогичен примеру 2, но время может быть на 20-30% меньше)

Пример 4: Распределение вычислений с DataParallel на 2 GPU

Модель копируется на все доступные GPU, батч делится поровну.

if torch.cuda.device_count() > 1:
    print(f'Используется {torch.cuda.device_count()} GPU')
    model = nn.DataParallel(model)
model.to(device)

# Цикл обучения без изменений
for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)
    outputs = model(images)
    loss = criterion(outputs, labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Используется 2 GPU
(обучение происходит быстрее, но нелинейно из-за накладных расходов)

Пример 5: Анализ использования памяти GPU с помощью memory_summary

Полезно для отладки утечек и оптимизации размера батча.

import torch

# Создадим большой тензор на GPU
t = torch.randn(1000, 1000, device='cuda')
print('После создания тензора:')
torch.cuda.memory_summary(device='cuda', abbreviated=True)

del t  # удалим ссылку
torch.cuda.empty_cache()
print('После очистки кэша:')
torch.cuda.memory_summary(device='cuda', abbreviated=True)

После создания тензора:
| alloc. requests   | 1                          |
| current memory    | 4.00 MB                    |
| peak memory       | 4.00 MB                    |
| ...
После очистки кэша:
| current memory    | 0.00 MB                    |

Пример 6: Использование CUDA streams для параллельного выполнения операций

Streams позволяют перекрыть передачу данных и вычисления.

s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

with torch.cuda.stream(s1):
    # Вычислительная операция A
    a = torch.randn(2000, 2000, device='cuda')
    b = torch.matmul(a, a)

with torch.cuda.stream(s2):
    # Операция B может выполняться одновременно, если ресурсы позволяют
    c = torch.randn(2000, 2000, device='cuda')
    d = c.t()

# Синхронизация
torch.cuda.synchronize()
print('Все потоки завершены')

Все потоки завершены

Пример 7: Использование pinned memory для ускорения загрузки данных

При передаче на GPU из CPU-памяти можно использовать закрепленные страницы (pinned memory), чтобы ускорить копирование.

from torch.utils.data import DataLoader, TensorDataset

# Используем pin_memory=True в DataLoader
dataloader = DataLoader(dataset, batch_size=64, pin_memory=True)

for batch in dataloader:
    # Данные уже могут быть в pinned memory, .to(device) работает быстрее
    batch = [t.to(device, non_blocking=True) for t in batch]

(Улучшение производительности на 5-15% в зависимости от размера данных)