如何在 PyTorch 框架中使用多个 GPU 进行推理

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（第一个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？我上面的脚本中有什么错误？

- AKSHAYAA VAIDYANATHAN

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

如果这个方法不起作用，请提供有关输入的更多细节。

[编辑]:

尝试这个：

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

- thedch

---

你打算如何使用批大小为1的多GPU推断？PyTorch应该如何在多个GPU之间分割数据？这不可能“开箱即用”。（它应该将模型分割到多个GPU上吗？应该将图像分成一半还是四分之一？） - thedch

请编辑您的问题，提供一个演示此功能的Keras代码片段。 - thedch

你能帮忙看一下这个问题吗？谢谢。https://stackoverflow.com/questions/74380417/how-to-use-multiple-gpu-at-the-same-time-when-using-models-with-pytorch-and-hugg - Furkan Gözükara