pytorch-如何从分布式数据并行学习中保存和加载模型

我不熟悉Pytorch DstributedDataParallel（），但我发现大多数教程都在培训期间保存了局部秩0模型。也就是说，如果我得到3台机器，每台机器上都有4个GPU，那么在最后的测试中，我会得到每台机器上保存的3个型号

例如，在第252行的pytorch教程中：

如果不是args.multiprocessing\u distributed或（args.multiprocessing\u distributed
和args.rank%ngpus_per_node==0）：
保存检查点（{…}）

如果

rank%ngpus\u per\u node==0

，则保存模型

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

那么为什么我们不把模型保存在

rank==0

，而

rank%ngpus\u per\u节点==0

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

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

如果我错了，请告诉我。

发生了什么事 如果我在任何地方出错，请纠正我

您所指的变更是在2018年通过引入的，描述如下：

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

以前，如果块，则保存的模型不带任何

，因此每个GPU上的每个节点都会保存一个确实浪费的模型，并且很可能会在每个节点上多次覆盖保存的模型

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

因此，每个进程的
args.rank
在脚本内部通过以下方式进行修改：

其评论如下：

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

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

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

如果您使用
rank==0
则每个世界只保存一个模型（其中世界将定义为
n个GPU*n个节点
）

问题
第一个问题
所以我们为什么不把模型保存在秩==0，但秩为%
每节点ngpus_==0

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

据我所知，DistributedDataParallel（）将自动运行
不做任何进一步的操作，就可以减少后端的损失
作业，每个进程都可以基于此自动同步丢失

准确地说，它与损失无关，而是根据文档（重点是我的文档）累积和应用权重修正：

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

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

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

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


假设保存模型的节点崩溃，文件丢失。你必须重做所有的东西。保存每个模型的操作成本不太高（每个历元执行一次或更少），因此可以轻松地为每个节点/工作者执行此操作
你必须重新开始训练。这意味着必须将模型复制到每个工作者（以及一些必要的元数据，尽管我认为这里不是这样）
节点无论如何都必须等待每个前向传递完成（因此可以平均梯度），如果模型保存需要大量时间，那么将浪费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

你读过《我是c》吗
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