pytorch - 如何从分布式数据并行学习中保存和加载模型
pytorch - How to Save and load model from DistributedDataParallel learning
我是 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 上保存模型？
如果我有多个模型，我应该使用哪个模型？
如果这是在分布式学习中保存模型的正确方法，我应该合并它们，使用其中一个，还是根据所有三个模型推断结果？
如果我错了请告诉我。
这是怎么回事
如有不妥之处请指正
您所指的更改是通过 this commit 在 2018 中引入的，描述为：
in multiprocessing mode, only one process will write the checkpoint
以前，这些是在 没有任何 if 块 的情况下保存的，因此每个 GPU 上的每个节点都会保存一个模型，这确实很浪费，并且很可能会覆盖多个已保存的模型每个节点上的时间。
现在，我们谈论的是分布式多处理（可能有很多工作人员，每个工作人员可能有多个 GPU）。
args.rank 对于每个进程，因此在脚本中由 this line:
修改
args.rank = args.rank * ngpus_per_node + gpu
其中有以下评论：
For multiprocessing distributed training, rank needs to be the
global rank among all the processes
因此 args.rank 是所有节点中所有 GPU 中的唯一 ID（或者看起来如此）。
如果是这样，并且每个节点都有 ngpus_per_node（在此训练代码中，假设每个节点都具有与我收集到的相同数量的 GPU），则模型仅保存一个（最后一个） ) 每个节点上的 GPU。在您使用 3 机器和 4 GPU 的示例中，您将获得 3 保存的模型（希望我能正确理解这段代码，因为它非常复杂 tbh）。
如果您使用 rank==0 每个世界 只有一个模型 （世界将被定义为 n_gpus * n_nodes）将被保存。
问题
第一题
So why don't we just save model on rank == 0, but rank %
ngpus_per_node == 0 ?
我将从你的假设开始，即：
To the best of my knowledge, DistributedDataParallel() will automatic
do all reduce to the loss on the backend, without doing any further
job, every process can sync the loss automatically base on that.
准确地说，它与损失无关，而是 gradient 累积和对权重应用修正，根据文档（强调我的）：
This container parallelizes the application of the given module by
splitting the input across the specified devices by chunking in the batch dimension. The module is replicated on each machine and
each device, and each such replica handles a portion of the input.
During the backwards pass, gradients from each node are averaged.
因此，当使用某些权重创建模型时，它会被复制到所有设备上（每个节点的每个 GPU）。现在每个 GPU 获得一部分输入（例如，对于等于 1024、4 个节点的总批大小，每个节点具有 4 个 GPU，每个 GPU 将获得 64 个元素），计算前向传递、损失，通过 .backward() 张量方法执行反向传播。现在所有梯度都通过 all-gather 平均，参数在 root 机器上优化，参数分布到所有节点，因此模块的状态在所有机器上始终相同。
注意：我不确定这种平均是如何发生的（而且我没有看到它在文档中明确说明），尽管我假设这些首先是平均的GPU 以及后来跨所有节点，因为我认为这将是最有效的。
现在，在这种情况下，为什么要为每个 node 保存模型？原则上你只能保存一个（因为所有模块将完全相同），但它有一些缺点：
- 假设保存模型的节点崩溃并且文件丢失。你必须重做所有的东西。保存每个模型的操作成本不是太高（每个 epoch 完成一次或更少）因此可以很容易地为每个 node/worker
完成
- 您必须重新开始训练。这意味着必须将模型复制到每个工作人员（以及一些必要的元数据，尽管我认为这里不是这种情况）
- 无论如何，节点都必须等待每次前向传递完成（因此可以平均梯度），如果模型保存花费大量时间，它将浪费 GPU/CPU 空闲（或其他一些同步）必须应用方案，我认为 PyTorch 中没有）。从整体上看，这有点“免费”。
问题 2（和 3）
And which model should I used for if I get multiple model?
没关系，因为所有这些都将完全相同，因为通过优化器对具有相同初始权重的模型应用相同的修正。
您可以使用这些方法来加载您保存的 .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
我是 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 上保存模型？
如果我有多个模型，我应该使用哪个模型？
如果这是在分布式学习中保存模型的正确方法，我应该合并它们，使用其中一个，还是根据所有三个模型推断结果？
如果我错了请告诉我。
这是怎么回事
如有不妥之处请指正
您所指的更改是通过 this commit 在 2018 中引入的，描述为：
in multiprocessing mode, only one process will write the checkpoint
以前，这些是在 没有任何 if 块 的情况下保存的，因此每个 GPU 上的每个节点都会保存一个模型，这确实很浪费，并且很可能会覆盖多个已保存的模型每个节点上的时间。
现在，我们谈论的是分布式多处理（可能有很多工作人员，每个工作人员可能有多个 GPU）。
args.rank 对于每个进程，因此在脚本中由 this line:
args.rank = args.rank * ngpus_per_node + gpu
其中有以下评论：
For multiprocessing distributed training, rank needs to be the global rank among all the processes
因此 args.rank 是所有节点中所有 GPU 中的唯一 ID（或者看起来如此）。
如果是这样，并且每个节点都有 ngpus_per_node（在此训练代码中，假设每个节点都具有与我收集到的相同数量的 GPU），则模型仅保存一个（最后一个） ) 每个节点上的 GPU。在您使用 3 机器和 4 GPU 的示例中，您将获得 3 保存的模型（希望我能正确理解这段代码，因为它非常复杂 tbh）。
如果您使用 rank==0 每个世界 只有一个模型 （世界将被定义为 n_gpus * n_nodes）将被保存。
问题
第一题
So why don't we just save model on rank == 0, but rank % ngpus_per_node == 0 ?
我将从你的假设开始，即：
To the best of my knowledge, DistributedDataParallel() will automatic do all reduce to the loss on the backend, without doing any further job, every process can sync the loss automatically base on that.
准确地说，它与损失无关，而是 gradient 累积和对权重应用修正，根据文档（强调我的）：
This container parallelizes the application of the given module by splitting the input across the specified devices by chunking in the batch dimension. The module is replicated on each machine and each device, and each such replica handles a portion of the input. During the backwards pass, gradients from each node are averaged.
因此，当使用某些权重创建模型时，它会被复制到所有设备上（每个节点的每个 GPU）。现在每个 GPU 获得一部分输入（例如，对于等于 1024、4 个节点的总批大小，每个节点具有 4 个 GPU，每个 GPU 将获得 64 个元素），计算前向传递、损失，通过 .backward() 张量方法执行反向传播。现在所有梯度都通过 all-gather 平均，参数在 root 机器上优化，参数分布到所有节点，因此模块的状态在所有机器上始终相同。
注意：我不确定这种平均是如何发生的（而且我没有看到它在文档中明确说明），尽管我假设这些首先是平均的GPU 以及后来跨所有节点，因为我认为这将是最有效的。
现在，在这种情况下，为什么要为每个 node 保存模型？原则上你只能保存一个（因为所有模块将完全相同），但它有一些缺点：
- 假设保存模型的节点崩溃并且文件丢失。你必须重做所有的东西。保存每个模型的操作成本不是太高（每个 epoch 完成一次或更少）因此可以很容易地为每个 node/worker 完成
- 您必须重新开始训练。这意味着必须将模型复制到每个工作人员（以及一些必要的元数据，尽管我认为这里不是这种情况）
- 无论如何，节点都必须等待每次前向传递完成（因此可以平均梯度），如果模型保存花费大量时间，它将浪费 GPU/CPU 空闲（或其他一些同步）必须应用方案，我认为 PyTorch 中没有）。从整体上看，这有点“免费”。
问题 2（和 3）
And which model should I used for if I get multiple model?
没关系，因为所有这些都将完全相同，因为通过优化器对具有相同初始权重的模型应用相同的修正。
您可以使用这些方法来加载您保存的 .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