可选: 数据并行

译者：@小王子

校对者：@李子文

作者: Sung Kim 和 Jenny Kang

在这个教程中, 我们将会学习如何在多个GPU上使用 DataParallel .

在 PyTorch 中使用 GPU 是一件很容易的事情.你可以像下面这样轻松的将一个模型分配到一个 GPU 上.

  1. model.gpu()

随后, 你可以将你的所有张量拷贝到上面的GPU:

  1. mytensor = my_tensor.gpu()

此处请注意: 如果只是调用 mytensor.gpu() 是不会将张量拷贝到 GPU 的.你需要将它赋给一个 新的张量, 这个张量就能在 GPU 上使用了.

在多个 GPU 上运行前向、反向传播是一件很自然的事情, 然而, PyTorch 默认情况下只会用到一个GPU, 可以通过使用 DataParallel 使你的模型并行运行, 在多个GPU上运行这些操作也将变得非常简单:

  1. model = nn.DataParallel(model)

这是教程的核心内容, 我们将在随后进行详细讲解

导入和参数

导入PyTorch模块和参数定义

  1. import torch
  2. import torch.nn as nn
  3. from torch.autograd import Variable
  4. from torch.utils.data import Dataset, DataLoader
  6. # 参数和数据加载
  7. input_size = 5
  8. output_size = 2
  10. batch_size = 30
  11. data_size = 100

伪数据集

只需要实现 getitem 就可以轻松的生成一个(随机）伪数据集, 如下代码所示:

  1. class RandomDataset(Dataset):
  3.     def __init__(self, size, length):
  4.         self.len = length
  5.         self.data = torch.randn(length, size)
  7.     def __getitem__(self, index):
  8.         return self.data[index]
  10.     def __len__(self):
  11.         return self.len
  13. rand_loader = DataLoader(dataset=RandomDataset(input_size, 100),
  14.                          batch_size=batch_size, shuffle=True)

简单模型

在下面的示例中, 我们的模型只需要一个输入并且完成一个线性操作, 最后得 到一个输出.当然, 你可以在任意模型 (CNN,RNN,Capsule Net等) 运用 DataParallel

我们在模型中设置了打印指令来监控输入和输出的张量大小, 请注意批数据次序为0时的输出.

  1. class Model(nn.Module):
  2.     # Our model
  4.     def __init__(self, input_size, output_size):
  5.         super(Model, self).__init__()
  6.         self.fc = nn.Linear(input_size, output_size)
  8.     def forward(self, input):
  9.         output = self.fc(input)
  10.         print("  In Model: input size", input.size(),
  11.               "output size", output.size())
  13.         return output

创建模型和 DataParallel

这是本教程的核心部分. 首先, 我们需要生成一个模型的实例并且检测我们是否拥有多个 GPU.如果有多个GPU , 我们可以使用 nn.DataParallel 来包装我们的模型, 然后我们 就可以将我们的模型通过 model.gpu() 施加于这些GPU上.

  1. model = Model(input_size, output_size)
  2. if torch.cuda.device_count() > 1:
  3.   print("Let's use", torch.cuda.device_count(), "GPUs!")
  4.   # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  5.   model = nn.DataParallel(model)
  7. if torch.cuda.is_available():
  8.    model.cuda()

运行模型

现在我们可以看到输入和输出张量的大小了.

  1. for data in rand_loader:
  2.     if torch.cuda.is_available():
  3.         input_var = Variable(data.cuda())
  4.     else:
  5.         input_var = Variable(data)
  7.     output = model(input_var)
  8.     print("Outside: input size", input_var.size(),
  9.           "output_size", output.size())

结果

当我们将输入设置为30批, 模型也产生了30批的输出.但是当我们使用多个GPU, 然后你 会得到类似下面这样的输出.

2 GPUs

如果有2个GPU, 我们将会看到这样的结果:

  1. # on 2 GPUs
  2. Let's use 2 GPUs!
  3.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  4.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  5. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  6.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  7.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  8. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  9.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  10.     In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
  11. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  12.     In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
  13.     In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
  14. Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

3 GPUs

如果有3个GPU, 我们将会看到这样的结果:

  1. Let's use 3 GPUs!
  2.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  3.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  4.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  5. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  6.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  7.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  8.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  9. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  10.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  11.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  12.     In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
  13. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  14.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  15.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  16.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  17. Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

8 GPUs

如果有8个GPU, 我们将会看到这样的结果:

  1. Let's use 8 GPUs!
  2.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  3.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  4.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  5.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  6.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  7.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  8.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  9.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  10. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  11.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  12.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  13.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  14.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  15.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  16.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  17.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  18.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  19. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  20.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  21.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  22.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  23.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  24.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  25.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  26.     In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
  27.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  28. Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
  29.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  30.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  31.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  32.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  33.     In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
  34. Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

总结

DataParallel 自动地将数据分割并且将任务送入多个GPU上的多个模型中进行处理. 在每个模型完成任务后, DataParallel 采集和合并所有结果, 并将最后的结果呈现给你.

想了解更多信息, 请点击: https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html.