DataParallel如何确定我要使用哪个GPU？

Question

我想找到一种简单的方法来指定实验运行的GPU。目前，我知道我可以在python命令前加上CUDA_VISIBLE_DEVICES = 1,2,3,4来设置gpu，并且我猜DataParallel随后将尝试使用所有gpu。

有没有一种方法可以直接告诉DataParallel id，例如4,7,9,12？

Answer 1

使用DataParallel，您可以使用以下语法指定要使用的设备：

model = torch.nn.DataParallel(model, device_ids=[ids_1,ids_2, ..., ids_n]).cuda()

使用CUDA_VISIBLE_DEVICES时，您正在设置代码可见的GPU。 出于种种原因，如果您设置CUDA_VISIBLE_DEVICES = 2,3然后执行：

model = torch.nn.DataParallel(model, device_ids=[0, 1]).cuda()

使用nvidia-smi，您会看到模型将在GPU 2和3上运行，因为您将特定ID设置为可见。

Answer 2

是的，DataParallem为我们提供了直接传递GPU ID的功能。

根据官方文档here，数据并行性是使用torch.nn.DataParallel实现的。可以将一个模块包装在DataParallel中，并且它将在批处理维度上通过多个GPU并行化。

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

对于您而言，您可以执行以下操作：

torch.nn.DataParallel(model, device_ids=[4, 7, 9, 12])
output = net(input_var)  # input_var can be on any device, including CPU

您可以在以下链接中了解有关如何将gpu id直接传递给DataParallel的更多信息：

• MULTI-GPU EXAMPLES
• DataParallel layers