如何在 pytorch 中使用多个 GPU？

Question

我正在学习pytorch，关注this tutorial。

我使用这个命令来使用 GPU。

device = torch.device("**cuda:0**" if torch.cuda.is_available() else "cpu")

但是，我想在 jupyter 中使用两个 GPU，如下所示：

device = torch.device("**cuda:0,1**" if torch.cuda.is_available() else "cpu")

当然，这是错误的。 那么，我该怎么做呢？

Answer 1

当我运行naiveinception_googlenet 时，上述方法对我不起作用。以下方法解决了我的问题。

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,3"  # specify which GPU(s) to be used

Answer 2

假设您希望在可用的 GPU 之间分配数据（如果您的批量大小为 16 和 2 个 GPU，您可能希望为每个 GPU 提供 8 个样本），而不是真正分散各个部分跨不同 GPU 的模型。这可以按如下方式完成：

如果您想使用所有可用的 GPU：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CreateModel()

model= nn.DataParallel(model)
model.to(device)

如果您想使用特定的 GPU： （例如，使用 4 个 GPU 中的 2 个）

device = torch.device("cuda:1,3" if torch.cuda.is_available() else "cpu") ## specify the GPU id's, GPU id's start from 0.

model = CreateModel()

model= nn.DataParallel(model,device_ids = [1, 3])
model.to(device)

通过设置操作系统环境变量来使用特定的 GPU：

在执行程序之前，设置CUDA_VISIBLE_DEVICES变量如下：

export CUDA_VISIBLE_DEVICES=1,3（假设您要选择第 2 和第 4 个 GPU）

然后，在程序中，您可以只使用DataParallel()，就好像您想使用所有 GPU。 （类似于第一种情况）。这里程序可用的 GPU 受操作系统环境变量的限制。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CreateModel()

model= nn.DataParallel(model)
model.to(device)

在所有这些情况下，数据都必须映射到设备。

如果X 和y 是数据：

X.to(device)
y.to(device)

Answer 3

另一种选择是为 PyTorch 使用一些帮助程序库：

PyTorch Ignite 库分布式 GPU 训练

在there 中有一个用于分布式配置的上下文管理器的概念：

• nccl - 多个 GPU 上的 Torch 原生分布式配置
• xla-tpu - TPU 分布式配置

PyTorch Lightning 多 GPU 训练

This 可能是恕我直言，在不更改原始 PyTorch 代码的情况下在 CPU/GPU/TPU 上进行训练的最佳选择。

对于类似的分布式 GPU 选项，值得一试 Catalyst。

Answer 4

使用多 GPU 就像在 DataParallel 中包装模型并增加批量大小一样简单。查看这两个教程以获得快速入门：

• Multi-GPU Examples
• Data Parallelism