如何在 pytorch 中使用多个 GPU? [英] How to use multiple GPUs in pytorch?

问题描述

我正在学习 pytorch 并遵循本教程.

I am learning pytorch and follow this tutorial.

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

I use this command to use a GPU.

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

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

But, I want to use two GPUs in jupyter, like this:

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

当然，这是错误的.那么，我该怎么做?

Of course, this is wrong. So, How can I do this?

推荐答案

假设您希望将数据分布在可用的 GPU 上(如果您的批次大小为 16 和 2 个 GPU，您可能希望提供 8 个样本到每个 GPU)，而不是真正将模型的各个部分分散到不同的 GPU 上.这可以按如下方式完成:

Assuming that you want to distribute the data across the available GPUs (If you have batch size of 16, and 2 GPUs, you might be looking providing the 8 samples to each of the GPUs), and not really spread out the parts of models across difference GPU's. This can be done as follows:

如果您想使用所有可用的 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 个)

If you want to use specific GPUs: (For example, using 2 out of 4 GPUs)

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变量如下:

Before executing the program, set CUDA_VISIBLE_DEVICES variable as follows:

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

export CUDA_VISIBLE_DEVICES=1,3 (Assuming you want to select 2nd and 4th GPU)

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

Then, within program, you can just use DataParallel() as though you want to use all the GPUs. (similar to 1st case). Here the GPUs available for the program is restricted by the OS environment variable.

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

model = CreateModel()

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

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

In all of these cases, the data has to be mapped to the device.

如果 X 和 y 是数据:

If X and y are the data:

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

这篇关于如何在 pytorch 中使用多个 GPU?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！