Как использовать multi-GPU во время вывода в рамках Pytorch

Question

Я пытаюсь сделать прогноз модели из unet3D, построенного на платформе pytorch.Я использую multi-gpus

import torch
import os
import torch.nn as nn
os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES']='0,1,2'

model = unet3d()
model = nn.DataParallel(model)
model = model.to('cuda')

result = model.forward(torch.tensor(input).to('cuda').float())

Но в модели по-прежнему используется только 1 графический процессор (первый), и я получаю ошибку памяти.

CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 11.00 GiB total capacity; 8.43 GiB already allocated; 52.21 MiB free; 5.17 MiB cached) 

Как я могу использовать несколько графических процессоровна этапе вывода?В чем ошибка в моем сценарии выше?

Answer 1

DataParallel обрабатывает отправку данных в GPU.

import torch
import os
import torch.nn as nn
os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES']='0,1,2'

model = unet3d()
model = nn.DataParallel(model.cuda())

result = model.forward(torch.tensor(input).float())

, если это не работает, пожалуйста, предоставьте более подробную информацию о input.

[EDIT]:

Попробуйте это:

with torch.no_grad():
    result = model(torch.tensor(input).float())