Question

我正在尝试从基于pytorch框架的unet3D进行模型预测。我正在使用多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)
model = model.to('cuda')

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

但是该模型仍然仅使用1个GPU（第一个），并且出现内存错误。

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)

在推理阶段如何使用多GPU？我上面的脚本中有什么错误？

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的更多详细信息。

[编辑]：

尝试一下：

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

pytorch框架中如何在推理过程中使用multi-gpu

1 个答案: