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

Posted 2023-03-16

参考技术A 本文源自 http://studyai.com/pytorch-1.4/beginner/blitz/data_parallel_tutorial.html

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

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

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

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

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

导入 PyTorch 模块和定义参数

设备

虚拟数据集

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

简单模型

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

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

创建模型和数据并行

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

运行模型

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

结果

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

总结

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