在Pytorch上使用多个GPU

Question

我有两个GPU，它们都是CUDA设备，但我不知道如何使用它们来训练我的resnext模型：GitHub存储库的link，我将其作为项目的基本代码。我更改的唯一重要的事情是：

    resnet152_model = resnet.resnext50_32x4d(pretrained=True)
    model = resnet152_model
    model = nn.DataParallel(model).to(device)

我的设备在哪里：

device = torch.device('cuda')

如果我编写cuda，它应该使用所有可用的GPU，但不是。它使用我的第一个GPU，如果我写，它将仅使用我的第二个GPU：

torch.cuda.set_device(1)

我的错误是什么以及如何使我的代码使用多个GPU

Answer 1

单身：

from torch import nn
import torch

# Set device to cuda if cuda is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Create model
model = nn.Sequential(
  nn.Conv2d(1,20,5),
  nn.ReLU(),
  nn.Conv2d(20,64,5),
  nn.ReLU()
)

# moving model to GPU
model.to(device)

对于多个：

# Optional DataParallel, not needed for single GPU usage
model1 = torch.nn.DataParallel(model, device_ids=[0]).to(device)
# Or, using default 'device_ids=None'
model1 = torch.nn.DataParallel(model).to(device)

parallelism_tutorial

如果要将训练数据集加载到多个GPU：

def model(data):
    for d in ['/gpu:0', '/gpu:1']:
        with tf.device(d):
            #You can do tf.nn.relu, max_pool_2x2, tf.nn.dropout, etc...

y = tf.nn.softmax(model(x))