如何在pytorch中使用多个GPU？

Question

我正在学习pytorch并遵循此教程：https ://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html

我使用此命令来使用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

假设您想将数据分布在可用的 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 2

使用多GPU就像包装模型DataParallel并增加批处理大小一样简单。查看以下两个教程以快速入门：

• 多GPU示例
• 数据并行