有没有办法在python多gpu环境中尴尬地并行运行程序？

Question

我有一个连接了4个GPU的节点。我有一个python代码，其中包含一个可以尴尬地并行化的循环。目前，我的程序仅使用1个GPU（我使用的库在GPU上运行dos模拟，但不支持多GPU）。 python中有没有一种方法可以在多个GPU上运行我的代码？我想要类似于下面的内容，但适用于GPU

from multiprocessing import Pool

def func(x):
    return x*x

if __name__ == '__main__':
    with Pool(5) as p:
        print(p.map(func, [1, 2, 3]))

Answer 1

我建议尝试一下Ray。参见GPU support。您基本上可以使用@ray.remote(num_gpus=1)来注释func(x)与func.remote(1)一起调用。