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)