一个 GPU 在训练期间比其他 GPU 使用更多的内存

Question

我使用 multigpu 用 pytorch 训练模型。一个 GPU 比其他 GPU 使用更多的内存，导致"内存不足"。为什么一个 GPU 会使用更多内存？是否有可能使使用更加平衡？还有其他方法可以减少内存使用量吗？(删除不再使用的变量...？批大小已为 1。谢谢。

Answer 1

DataParallel拆分批处理并将每个拆分发送到不同的 GPU，每个 GPU 都有一个模型副本，然后独立计算前向传递，然后将每个 GPU 的输出收集回一个 GPU，而不是在每个 GPU 中独立计算损失。

如果要缓解此问题，可以在DataParallel模块中包含损失计算。

如果这样做仍然是一个问题，那么您可能需要模型并行性而不是数据并行性：使用.cuda(gpu_id)将模型的不同部分移动到不同的 GPU。当模型的权重非常大时，这很有用。