PyTorch 在第一个可用的 GPU (cuda:0) 上分配更多内存答案

【问题标题】：PyTorch allocates more memory on the first available GPU (cuda:0)PyTorch 在第一个可用的 GPU (cuda:0) 上分配更多内存
【发布时间】：2020-05-09 10:20:51
【问题描述】：

作为强化学习训练系统的一部分，我使用四个 GPU 并行训练四个策略。对于每个模型，有两个进程——参与者和学习者，它们只使用它们特定的 GPU（例如，对应于模型 #2 的参与者和学习者只使用 GPU #2 来处理它们的所有张量）。 Actor 和Learner 通过torch 的share_memory_() 共享模型层。由于四个训练“子系统”是完全对称的，我希望它们在四个 GPU 上使用完全相同数量的 GPU 内存。然而，在实践中，我发现在第一个 GPU (cuda:0) 上分配了更多的 GPU 内存。

感觉所有的内存共享都是通过 GPU #0 完成的。有没有办法解决这个问题？

到目前为止，我尝试通过显式更改进程start 函数中的os.environ 在子进程中设置CUDA_VISIBLE_DEVICES。这似乎没有任何效果，可能是因为子进程是从主进程派生出来的，其中 PyTorch CUDA 已经初始化，并且此时似乎忽略了 envvars。

【问题讨论】：

如果我有能力拥有四个 2080 RTX，那么我很乐意尝试并提供帮助。这是我很想遇到的问题。
公司为此类事情付费；）

标签： deep-learning gpu pytorch

【解决方案1】：

好的，到目前为止，我想出了一个解决方法。我的假设是正确的，如果 PyTorch CUDA 子系统在子进程被派生之前已经初始化，那么将 CUDA_VISIBLE_DEVICES 设置为子进程的不同值不会做任何事情。

更糟糕的是，调用torch.cuda.device_count() 足以初始化 CUDA，所以我们甚至无法从 PyTorch 查询 GPU 的数量。解决方案是硬编码，作为参数传递，或者在单独的进程中查询 PyTorch API。我对后者的实现：

import sys


def get_available_gpus_without_triggering_pytorch_cuda_initialization(envvars):
    import subprocess
    out = subprocess.run([sys.executable, '-m', 'utils.get_available_gpus'], capture_output=True, env=envvars)
    text_output = out.stdout.decode()
    from utils.utils import log
    log.debug('Queried available GPUs: %s', text_output)
    return text_output


def main():
    import torch
    device_count = torch.cuda.device_count()
    available_gpus = ','.join(str(g) for g in range(device_count))
    print(available_gpus)
    return 0


if __name__ == '__main__':
    sys.exit(main())

基本上这个函数调用它自己的脚本作为单独的python进程并读取stdout。

我不会将此答案标记为已接受，因为我想了解一个合适的解决方案（如果存在）。

【讨论】：