Проблемы с torch.nn.DataParallel

Question

Я новичок в области глубокого обучения.Сейчас я воспроизводлю бумажные коды.поскольку они используют несколько графических процессоров, в кодах есть команда torch.nn.DataParallel(model, device_ids= args.gpus).cuda().Но у меня есть только один графический процессор, что я должен изменить этот код, чтобы соответствовать моему графическому процессору?

Спасибо!

Answer 1

DataParallel также должен работать на одном GPU, но вы должны проверить, содержит ли args.gpus только идентификатор устройства, которое должно использоваться (должно быть 0) или None.Выбор None заставит модуль использовать все доступные устройства.

Также вы можете удалить DataParallel, так как он вам не нужен, и переместить модель в графический процессор только путем вызова model.cuda() или, как я предпочитаю,model.to(device) где device - имя устройства.

Пример:

В этом примере показано, как использовать модель на одном графическом процессоре, настраивая устройство с помощью .to() вместо .cuda().

from torch import nn
import torch

# Set device to cuda if cuda is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Create model
model = nn.Sequential(
  nn.Conv2d(1,20,5),
  nn.ReLU(),
  nn.Conv2d(20,64,5),
  nn.ReLU()
)

# moving model to GPU
model.to(device)

Если вы хотите использовать DataParallel, вы можете сделать это следующим образом

# Optional DataParallel, not needed for single GPU usage
model1 = torch.nn.DataParallel(model, device_ids=[0]).to(device)
# Or, using default 'device_ids=None'
model1 = torch.nn.DataParallel(model).to(device)