Deep learning PyTorch在第一个可用GPU上分配更多内存（cuda:0）

作为强化学习培训系统的一部分，我正在使用四个GPU并行培训四个策略。对于每个模型，有两个过程-参与者和学习者，它们仅使用其特定的GPU（例如，对应于模型2的参与者和学习者仅对其所有张量使用GPU 2）。演员和学习者通过torch的

共享内存（）

共享模型层。 由于四个训练“子系统”是完全对称的，我希望它们在四个GPU上使用完全相同的GPU内存量。然而，实际上，我看到在第一个GPU上分配了更多的GPU内存（

cuda:0

）

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

到目前为止，我试图通过在进程

start

函数中显式更改

os.environ

在子进程中设置

CUDA\u VISIBLE\u DEVICES

。这似乎没有任何影响，可能是因为子进程是从主进程派生出来的，在主进程中PyTorch CUDA已经初始化，而envvars似乎在这一点上被忽略了

---

好的，到目前为止，我想出了一个解决办法。我的假设是正确的，如果PyTorch CUDA子系统在子进程分叉之前已经初始化，那么为子进程将CUDA_VISIBLE_设备设置为不同的值不会起任何作用

更糟糕的是，调用

torch.cuda.device\u count（）

就足以初始化cuda，因此我们甚至无法从PyTorch查询GPU的数量。解决方案要么是硬编码，作为参数传递，要么是在单独的过程中查询pytorchapi。我对后者的实施：

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

---

我不会将此答案标记为已接受，因为如果存在，我想学习适当的解决方案。

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