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

Posted 2023-03-28

【中文标题】pytorch - 如何从 DistributedDataParallel 学习中保存和加载模型

【英文标题】：pytorch - How to Save and load model from DistributedDataParallel learning

【发布时间】：2020-07-17 03:41:25

【问题描述】：

我是 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？

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

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

如果我错了，请告诉我。

【问题讨论】：

你读过pytorch.org/tutorials/intermediate/…吗？我很好奇，哪一部分对你来说没有意义？ （顺便说一句，我发现它也是一个非常令人困惑的教程）

@CharlieParker 教程很清楚。但从 ImageNet 示例代码中，他们为每台计算机保存了模型。所以我只是想知道从每台计算机保存的每个模型之间是否有任何不同。训练后我应该使用哪一个。 Szymon Maszke 的回答很棒。我明白了。

【参考方案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