PyTorch 中的多 GPU 训练和梯度累积作为替代方案
发布网友
发布时间：2024-10-04 19:21
共1个回答
热心网友
时间：2024-10-05 01:36
在本文中，我们将探讨数据并行（DP）和分布式数据并行（DDP）算法的差异，并解释梯度累积（GA）的概念。此外，我们还将展示DDP和GA在PyTorch中的实现方式及其如何导致相同的结果。
在训练深度神经网络（DNN）时，批量大小是一个关键的超参数。批量大小不宜过大，以免网络过拟合；也不宜过小，以免收敛速度慢。
当处理高分辨率图像或其他大型数据时，我们通常在GPU上进行DNN模型的训练。根据GPU内存的大小，拟合小批量大小可能会出现问题。为了增加有效批量大小，我们可以采取以下三种主要方法：
1. 如果您拥有大型GPU，可以在其上容纳所需的所有数据，您可以跳过以下部分，直接阅读DDP部分，并在完整代码部分中查看其在PyTorch中的实现。
2. 当希望有效批量大小为30，但每个GPU上只能容纳10个数据点时，我们可以选择数据并行或分布式数据并行。
3. 如果只有一个GPU，但希望使用更大的批量大小，另一种选择是累积一定数量的步骤的梯度，从而增加有效的批量大小。
数据并行性（DP）：
首先，我们定义主GPU。然后，执行以下步骤：
这个过程存在一些问题和低效率。
分布式数据并行（DDP）：
引入分布式数据并行是为了改善数据并行算法的低效率。我们仍然采用与之前相同的设置——每批30个数据点，使用3个GPU。差异如下：
由于DDP相对于DP的优点，目前优先使用DDP，因此我们只展示DDP的实现。
梯度累积：
如果我们只有一个GPU但想使用更大的批量大小，另一种选择是累积一定数量的步骤的梯度。从上面的例子中，我们可以通过3次迭代累积10个数据点的梯度，以达到与我们在有效批量大小为30的DDP训练中描述的结果相同的结果。
DDP流程代码：
下面我将仅介绍与1 GPU代码相比实现DDP时的差异。完整的代码可以在以下部分找到。首先，初始化进程组，允许不同进程之间进行通信。使用int(os.environ[“LOCAL_RANK”])检索给定进程中使用的GPU。然后，将模型包装在DistributedDataParallel中，以支持多GPU训练。最后一部分是定义我在DDP部分中提到的DistributedSampler。培训的其余部分保持不变，将在本文末尾包含完整的代码。
梯度累积代码：
当反向传播发生时，在我们调用loss.backward()后，梯度将存储在各自的张量中。实际的更新发生在调用optimizer.step()时，然后使用optimizer.zero_grad()将张量中存储的梯度设置为零，以运行反向传播和参数更新的下一次迭代。为了累积梯度，我们调用loss.backward()获取所需的梯度累积数量，而不将梯度设置为零，以便它们在多次迭代中累积，然后对它们进行平均以获得累积梯度迭代中的平均梯度（loss = loss/ACC_STEPS）。之后，调用optimizer.step()并将梯度归零以开始下一次梯度累积。
代码总结：
在本文中，我们简要介绍并直观地介绍了DP、DDP算法和梯度累积，并展示了如何在没有多个GPU的情况下增加有效批量大小。需要注意的一件重要事情是，即使我们获得相同的最终结果，使用多个GPU进行训练也比使用梯度累积要快得多，因此如果训练速度很重要，那么使用多个GPU是加速训练的唯一方法。
