PyTorch等深度學習框架在使用多個GPU時如何處理內存？

Question

我最近遇到了一個單個Nvidia V100內存不足的情況。 我使用多個GPU來訓練網絡的經驗有限，所以我對數據並行化過程的工作方式有些不確定。 假設我使用的是模型和批量大小，需要20-25GB的內存。 有沒有辦法充分利用兩個16GB V100之間的全部32GB內存？ PyTorch的DataParallel功能會實現嗎？ 我想也有可能打破模型並使用模型並行性。 請原諒我對這個問題缺乏了解。 在此先感謝您的幫助或澄清！

Answer 1

你應該保持模型並行性作為你的最后一個資源，並且只有你的模型不適合單個GPU的內存（16GB / GPU你有足夠的空間容納一個巨大的模型）。

如果你有兩個GPU，我會使用數據並行。 在數據並行性方面，您可以在每個GPU上獲得模型的副本，並為每個副本提供批處理。 然后收集漸變並用於更新副本。

Pytorch使得實現數據並行化變得非常容易，因為您只需要在nn.DataParallel包裝模型實例：

model = torch.nn.DataParallel(model, device_ids=[0, 1])
output = model(input_var)