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

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

---

是否有一种方法可以直接告诉DataParallel ID，如4、7、9、12？

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

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

当您使用CUDA_VISIBLE_设备时，您将GPU设置为代码可见。 对于iNtance，如果将CUDA\u VISIBLE\u DEVICES设置为2,3，然后执行：

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

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

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

根据官方文件，使用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的更多信息：