代码之家 › 专栏 › 技术社区 › MenorcanOrange

pytorch框架下如何使用多gpu进行推理

multi-gpu pytorch

0

MenorcanOrange · 技术社区 · 5 年前

我正在尝试从构建在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(第一个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?我上面的剧本有什么错误?

0 回复 | 直到 5 年前

1

3

thedch 5 年前

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())