如何使用Python在多个进程中运行pycuda

Question

我有一个可以在单个进程中运行的 pycuda 代码。python的多进程能否支持在多个子进程中运行这段代码？如果我尝试，我会发现我犯了一个错误。我犯错了吗？尝试用python的进程实现一个简单的多进程，发现会出错。    import pycuda.autoinit    import pycuda.driver as drv    import numpy    from pycuda.compiler import SourceModule    from multiprocessing import Pool, Manager, Process    def ffunc(i, return_dict, a, b, multiply_them):        dest = numpy.zeros_like(a)        multiply_them(            drv.Out(dest), drv.In(a), drv.In(b),            block=(400, 1, 1), grid=(1, 1))        return_dict[i] = dest    if __name__ == '__main__':        mod = SourceModule("""        __global__ void multiply_them(float *dest, float *a, float *b)        {         const int i = threadIdx.x;         dest[i] = a[i] * b[i];        }        """)        multiply_them = mod.get_function("multiply_them")        aa = numpy.random.randn(2, 400).astype(numpy.float32)        bb = numpy.random.randn(2, 400).astype(numpy.float32)        manager = Manager()        return_dict = manager.dict()        jobs = []        for i in range(2):            p = Process(target=ffunc, args=(i, return_dict, aa[i], bb[i], multiply_them))            jobs.append(p)            p.start()        for p in jobs:            p.join()        print(return_dict)我不确定 pycuda 是否可以在不同的进程中运行。我期待您的建议。

Answer 1

幸运的是，我解决了这个问题。

在main函数中添加一行代码：

multiprocessing.set_start_method('spawn')

Answer 2

CUDA 不应该在分叉之前初始化。

您可以在此处找到更多详细信息：https ://forums.developer.nvidia.com/t/cuda8-0-bug-child-process-forked-after-cuinit-get-cuda-error-not-initialized-on-cuinit /45764

最好产生一个新进程；因此，multiprocessing.set_start_method('spawn')完美运行。