GPU-обучение

Переносим полноценное обучение на GPU и разбираем подводные камни устройств.

GPU-обучение — это то же обучение, но модель и каждый батч данных переносятся на видеокарту через .to(device).

Три места, где появляется device

В разделе про тензоры мы выбрали device и научились переносить тензоры. Теперь применим это к обучению. Чтобы считать на GPU, на устройство нужно перенести три вещи: модель и каждый батч (входы и метки). Правило железное: модель и данные обязаны быть на одном устройстве, иначе ошибка.

import torch

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

# 1. модель — один раз
model = Net().to(device)

for epoch in range(num_epochs):
    for X_batch, y_batch in loader:
        # 2. данные — каждый батч
        X_batch = X_batch.to(device)
        y_batch = y_batch.to(device)

        pred = model(X_batch)
        loss = criterion(pred, y_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Модель переносят один раз до цикла (её веса остаются на GPU всё обучение). А батчи — каждый раз внутри цикла, потому что DataLoader отдаёт их с CPU. Это не лишняя работа: данные физически лежат в обычной памяти, и каждый новый батч нужно доставить на видеокарту.

Почему именно так, а не разом

Возникает соблазн перенести весь датасет на GPU сразу. Иногда для маленьких данных так и делают, но обычно весь датасет в память видеокарты не влезает — её там немного. Поэтому переносят порциями-батчами, ровно столько, сколько нужно для текущего шага.

Частые ошибки устройств

Ошибка	Причина
Expected all tensors on same device	модель на GPU, а батч забыли перенести
can't convert cuda tensor to numpy	забыли `.cpu()` перед `.numpy()`
метки и предсказания на разных устройствах	перенесли только `X`, забыли `y`

Самая частая — перенесли входы, но забыли метки y. Переносите оба тензора батча. И помните формулу вывода результата с GPP в numpy из второго раздела: tensor.detach().cpu().numpy().

Когда GPU реально помогает

GPU ускоряет за счёт массового параллелизма матричных операций. Выгода тем больше, чем больше модель и батч: на крошечной линейной регрессии разницы с CPU почти нет (накладные расходы на перенос съедают выигрыш), а на большой свёрточной сети ускорение — десятки раз. Поэтому не удивляйтесь, если учебный пример на GPU не станет быстрее: эффект проявляется на серьёзных моделях.

Сохранение и устройства

Полезная деталь: при загрузке весов можно сразу указать, куда их класть, через аргумент map_location. Это спасает, когда модель обучали на GPU, а грузят на машине только с CPU:

# обучали на GPU, грузим на машине без GPU
state = torch.load("model.pth", map_location="cpu")
model.load_state_dict(state)

Итог

Для GPU переносят на device: модель (один раз) и каждый батч входов и меток (в цикле).
Модель и данные обязаны быть на одном устройстве, иначе ошибка.
Самая частая ошибка — перенесли X, но забыли y.
Выигрыш GPU заметен на больших моделях и батчах; на учебных примерах его может не быть.