如何在pytorch框架的推理过程中使用多GPU

【问题标题】：How to use multi-gpu during inference in pytorch framework如何在pytorch框架的推理过程中使用多GPU
【发布时间】：2019-07-13 05:59:16
【问题描述】：

我正在尝试从基于 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？我上面的脚本有什么错误？

【问题讨论】：

标签： pytorch multi-gpu

【解决方案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())

【讨论】：

您打算如何以 1 的批大小进行多 GPU 推理？ pytorch 应该如何跨 GPU 拆分数据？这是不可能的“开箱即用”。（它应该在 GPU 上拆分模型吗？它应该将图像分成两半还是四分之一？）。
通过将权重和层分布在不同的 gpus 上，在 keras 中也会发生同样的事情
请编辑您的问题以提供演示此功能的 Keras sn-p。