如何使用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')完美运行。