pytorch - 如何从 DistributedDataParallel 学习中保存和加载模型

Question

我是 Pytorch DstributedDataParallel() 的新手，但我发现大多数教程在训练期间保存了 local rank 0 模型。这意味着，如果我得到 3 台机器，每台机器配备 4 个 GPU，那么在决赛中，我将得到 3 个模型，每台机器都可以节省。

例如在 pytorch ImageNet 教程第 252 行：

if not args.multiprocessing_distributed or (args.multiprocessing_distributed
                and args.rank % ngpus_per_node == 0):
            save_checkpoint({...})

如果rank % ngpus_per_node == 0，他们会保存模型。

据我所知，DistributedDataParallel() 会自动减少后端的损失，无需做任何进一步的工作，每个进程都可以在此基础上自动同步损失。 每个流程上的所有模型只会在流程结束时略有不同。这意味着我们只需要保存一个模型就足够了。

那么我们为什么不把模型保存在rank == 0，而是rank % ngpus_per_node == 0？

如果我有多个模型，我应该使用哪个模型？

如果这是在分布式学习中保存模型的正确方法，我应该合并它们，使用其中一个，还是根据所有三个模型推断结果？

如果我错了，请告诉我。

Answer 1

发生了什么

如有错误请指正

您所指的更改是通过this commit 在2018 中引入的，并描述为：

在多进程模式下，只有一个进程会写入检查点

以前，那些保存没有任何if 块，因此每个 GPU 上的每个节点都会保存一个模型，这确实很浪费，并且很可能会在每个节点上多次覆盖保存的模型。 p>

现在，我们讨论的是分布式多处理（可能有很多工作人员，每个工作人员可能有多个 GPU）。

每个进程的args.rank因此在脚本内部由this line修改：

args.rank = args.rank * ngpus_per_node + gpu

有以下评论：

对于多处理分布式训练，排名需要是 所有进程中的全局排名

因此args.rank 在所有节点的所有 GPU 中是唯一 ID（或者看起来如此）。

如果是这样，并且每个节点都有 ngpus_per_node（在此训练代码中，假设每个节点都具有我收集的相同数量的 GPU），那么模型仅保存用于每个节点上的一个（最后一个）GPU节点。在您使用3 机器和4 GPU 的示例中，您将获得3 保存的模型（希望我正确理解此代码，因为它非常复杂）。

如果您使用rank==0，每个世界（世界将被定义为n_gpus * n_nodes）只会保存一个模型。

问题

第一个问题

那么我们为什么不把模型保存在 rank == 0 上，但是 rank % ngpus_per_node == 0 ?

我将从你的假设开始，即：

据我所知，DistributedDataParallel() 将自动 尽一切努力减少后端的损失，无需做任何进一步的事情 作业，每个进程都可以在此基础上自动同步丢失。

确切地说，它与损失无关，而是 gradient 累积和对权重应用更正，根据文档（强调我的）：

此容器通过以下方式并行化给定模块的应用程序 在指定的设备上拆分输入，通过在批处理维度中进行分块。 模块在每台机器上复制，并且 每个设备，每个这样的副本处理输入的一部分。 在向后传递期间，每个节点的梯度被平均。

因此，当使用某些权重创建模型时，它会在所有设备（每个节点的每个 GPU）上复制。现在每个 GPU 获得一部分输入（例如，对于等于 1024、4 节点的总批量大小，每个节点都有 4 GPU，每个 GPU 将获得 64 元素），计算前向传递、损失、执行反向传播通过.backward()张量方法。现在所有梯度均由 all-gather 平均，参数在 root 机器上进行了优化，参数分布到所有节点，因此模块的状态在所有机器上始终相同。

注意：我不确定这种平均是如何发生的（我没有在文档中明确说明），尽管我假设这些首先是跨 GPU 平均，然后是所有节点，因为这将是我认为最有效的。

现在，在这种情况下，为什么要为每个 node 保存模型？原则上你只能保存一个（因为所有模块都完全相同），但它有一些缺点：

• 说您保存模型的节点崩溃并且文件丢失。你必须重做所有的东西。保存每个模型不是太昂贵的操作（每个 epoch 完成一次或更少），因此可以为每个节点/worker 轻松完成
• 您必须重新开始训练。这意味着必须将模型复制到每个工作人员（以及一些必要的元数据，尽管我认为这里不是这种情况）
• 节点将不得不等待每个前向传递完成（因此梯度可以被平均），如果模型保存需要很长时间，它会浪费 GPU/CPU 空闲（或者其他一些同步方案必须被应用，我认为 PyTorch 中没有）。如果您从整体上看，这有点“免费”。

问题 2（和 3）

如果我有多个模型，我应该使用哪个模型？

这并不重要，因为它们都将完全相同，因为通过优化器将相同的校正应用于具有相同初始权重的模型。

您可以使用类似的方法来加载您保存的 .pth 模型：

import torch

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

parallel_model = torch.nn.DataParallel(MyModelGoesHere())
parallel_model.load_state_dict(
    torch.load("my_saved_model_state_dict.pth", map_location=str(device))
)

# DataParallel has model as an attribute
usable_model = parallel_model.model