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

Question

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

有没有办法直接告诉 DataParallel id，比如 4、7、9、12？

标签： pytorch

Answer 1

是的，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

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

• 多 GPU 示例
• 数据并行层