PyTorch-4 nn.DataParallel 数据并行详解
发布网友
发布时间：2022-11-11 07:37
共1个回答
热心网友
时间：2024-12-13 01:53
在本教程中，我们将学习如何使用多个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 在将结果返回给你之前收集和合并结果。
