pytorch - 如何從分布式數據並行學習中保存和加載 model

Question

我是 Pytorch DstributedDataParallel() 的新手，但我發現大多數教程在訓練期間都保存了本地排名 0 model。 這意味着如果我得到 3 台機器，每台機器上都有 4 個 GPU，最后我會得到 3 個 model，從每台機器上保存。

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

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

如果rank % ngpus_per_node == 0 ，他們會保存 model 。

據我所知，DistributedDataParallel() 將自動減少后端的損失，無需做任何進一步的工作，每個進程都可以在此基礎上自動同步損失。 每個進程上的所有 model 只會在進程結束時略有不同。 這意味着我們只需要保存一個 model 就足夠了。

那么為什么我們不把 model 保存在rank == 0上，但是rank % ngpus_per_node == 0呢？

如果我得到多個 model，我應該使用哪個 model？

如果這是在分布式學習中保存 model 的正確方法，我應該合並它們，使用其中一個，還是根據所有三個模型推斷結果？

如果我錯了，請告訴我。

Answer 1

到底是怎么回事

如果我在任何地方錯了，請糾正我

您所指的更改是在2018通過此提交引入的，並描述為：

在多處理模式下，只有一個進程會寫入檢查點

以前，這些是在沒有任何if塊的情況下保存的，因此每個 GPU 上的每個節點都會保存一個 model，這確實很浪費，並且很可能會在每個節點上多次覆蓋保存的 model。

現在，我們正在討論分布式多處理（可能有很多工作人員，每個工作人員可能有多個 GPU）。

因此，每個進程的args.rank在腳本內通過以下行進行修改：

args.rank = args.rank * ngpus_per_node + gpu

其中有以下評論：

對於多進程分布式訓練，rank 需要是所有進程中的全局排名

因此args.rank是所有節點中所有 GPU 中的唯一 ID （或者看起來如此）。

如果是這樣，並且每個節點都有ngpus_per_node （在此訓練代碼中，假設每個節點都具有與我收集的相同數量的 GPU），則 model 僅保存每個節點上的一個（最后一個）GPU。 在您使用3台機器和4 GPU 的示例中，您將獲得3已保存的模型（希望我正確理解此代碼，因為它非常復雜）。

如果您使用rank==0 ，則每個世界（其中世界將被定義為n_gpus * n_nodes ）只會保存一個 model 。

問題

第一個問題

那么為什么我們不把 model 保存在 rank == 0 上，但是 rank % ngpus_per_node == 0 呢？

我將從您的假設開始，即：

據我所知，DistributedDataParallel() 將自動減少后端的損失，無需做任何進一步的工作，每個進程都可以在此基礎上自動同步損失。

准確地說，它與損失無關，而是gradient累積和對權重的應用校正，根據文檔（強調我的）：

此容器通過在批處理維度中分塊將輸入拆分到指定的設備，從而並行化給定模塊的應用程序。 該模塊在每台機器和每台設備上復制，每個這樣的副本處理輸入的一部分。 在向后傳遞期間，來自每個節點的梯度被平均。

因此，當使用某些權重創建 model 時，它會在所有設備上復制（每個節點的每個 GPU）。 現在每個 GPU 獲得一部分輸入（例如，對於等於1024的總批大小， 4個節點，每個節點有4 GPU，每個 GPU 將獲得64元素），計算前向傳遞，損失，通過.backward()張量方法執行反向傳播。 現在所有梯度均由 all-gather 平均，參數在root機器上優化，參數分布到所有節點，因此模塊的 state 在所有機器上始終相同。

注意：我不確定這種平均是如何發生的（我沒有在文檔中明確說明），盡管我假設這些首先在 GPU 上平均，然后在所有節點上平均，因為這將是我認為最有效的.

現在，在這種情況下，為什么要為每個node保存 model 呢？ 原則上你只能保存一個（因為所有模塊都完全相同），但它有一些缺點：

• 假設您保存 model 的節點崩潰並且文件丟失。 你必須重做所有的東西。 保存每個 model 的操作成本不會太高（每個 epoch 完成一次或更少），因此每個節點/worker 都可以輕松完成
• 你必須重新開始訓練。 這意味着必須將 model 復制到每個工作人員（以及一些必要的元數據，盡管我認為這里不是這種情況）
• 無論如何，節點都必須等待每個前向傳遞完成（因此可以平均梯度），如果 model 保存需要大量時間，它將浪費 GPU/CPU 空閑（或必須應用一些其他同步方案，我認為 PyTorch 中沒有）。 如果您查看整體情況，這使得它有點“免費”。

問題 2（和 3）

如果我得到多個 model，我應該使用哪個 model？

沒關系，因為它們都將完全相同，因為通過優化器將相同的校正應用於具有相同初始權重的 model。

您可以使用這些方法來加載您保存的.pth model：

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