PyTorch-4 nn.DataParallel 数据并行详解

最新推荐文章于 2024-07-07 16:25:33 发布

ScorpioDoctor

最新推荐文章于 2024-07-07 16:25:33 发布

阅读量1.4w

点赞数 15

分类专栏： PyTorch教程文章标签： pytorch

本文链接：https://blog.csdn.net/zhjm07054115/article/details/104799661

版权

PyTorch教程专栏收录该内容

6 篇文章 2 订阅

订阅专栏

本文源自 http://studyai.com/pytorch-1.4/beginner/blitz/data_parallel_tutorial.html

在本教程中，我们将学习如何使用多个GPU： DataParallel 的用法.

与PyTorch一起使用GPU非常容易。您可以将模型放在GPU上:

device = torch.device("cuda:0")
model.to(device)

Then, you can copy all your tensors to the GPU:

mytensor = my_tensor.to(device)

请注意，只要调用 my_tensor.to(device) ，就会在GPU上返回 my_tensor 的新副本，而不是重写 my_tensor 。您需要将它分配给一个新的tensor，并在GPU上使用该tensor。

在多个GPU上执行前向、后向传播是很自然的。但是，PYTORCH默认只使用一个GPU。你可以轻松地在多个GPU上运行您的操作，方法是让你的模型使用 DataParallel 并行运行:

model = nn.DataParallel(model)

这是本教程的核心。我们将在下面更详细地探讨它。导入与参数 ———————-

导入 PyTorch 模块和定义参数

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

参数和数据加载器 DataLoaders

input_size = 5
output_size = 2

batch_size = 30
data_size = 100

设备

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

虚拟数据集

制造一个虚拟的(随机产生) 数据集。你只需要实现 Python 的魔法函数 getitem :

class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len

rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
                         batch_size=batch_size, shuffle=True)

简单模型

对于演示，我们的模型只获得一个输入，执行一个线性操作，并给出一个输出。但是，您可以在任何模型(CNN、RNN、Capsule Net等)上使用 DataParallel 。

我们在模型中放置了一个print语句来监视输入和输出张量的大小。请注意批次0的打印内容。

class Model(nn.Module):
    # 我们的模型

    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(),
              "output size", output.size())

        return output

创建模型和数据并行

这是本教程的核心部分。首先，我们需要创建一个模型实例，并检查我们是否有多个GPU。如果我们有多个GPU，我们可以使用 nn.DataParallel 来包装我们的模型。然后我们可以通过模型 model.to(device) 将我们的模型放在GPU上

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

model.to(device)

运行模型

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

for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("Outside: input size", input.size(),
          "output_size", output.size())

结果

如果您没有GPU或一个GPU，当我们批处理30个输入和30个输出时，模型得到30，输出与预期相同。但是如果你有多个GPU，那么你可以得到这样的结果。

# 2 GPUs
# ~~~~~~
#
# If you have 2, you will see:
#
# .. code:: bash
#
#     # on 2 GPUs
#     Let's use 2 GPUs!
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
#         In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
#         In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2])
#     Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
#
# 3 GPUs
# ~~~~~~
#
# If you have 3 GPUs, you will see:
#
# .. code:: bash
#
#     Let's use 3 GPUs!
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#         In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#     Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
#
# 8 GPUs
# ~~~~~~~~~~~~~~
#
# If you have 8, you will see:
#
# .. code:: bash
#
#     Let's use 8 GPUs!
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([4, 5]) output size torch.Size([4, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#     Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#         In Model: input size torch.Size([2, 5]) output size torch.Size([2, 2])
#     Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])
#

总结

DataParallel 会自动拆分数据，并将作业订单发送到多个GPU上的多个模型。在每个模型完成它们的工作之后，DataParallel 在将结果返回给你之前收集和合并结果。