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

我有一个可以在单个进程中运行的 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)

Process Process-2:
Traceback (most recent call last):
  File "/home/vision/anaconda3/envs/py3b/lib/python3.6/multiprocessing/process.py", line 258, in _bootstrap
    self.run()
  File "/home/vision/anaconda3/envs/py3b/lib/python3.6/multiprocessing/process.py", line 93, in run
    self._target(*self._args, **self._kwargs)
  File "/home/vision/lpx/AE23D/test_pycuda.py", line 22, in ffunc
    block=(400,1,1), grid=(1,1))
  File "/home/vision/anaconda3/envs/py3b/lib/python3.6/site-packages/pycuda/driver.py", line 382, in function_call
    func._set_block_shape(*block)
pycuda._driver.LogicError: cuFuncSetBlockShape failed: initialization error
Process Process-3:
Traceback (most recent call last):
  File "/home/vision/anaconda3/envs/py3b/lib/python3.6/multiprocessing/process.py", line 258, in _bootstrap
    self.run()
  File "/home/vision/anaconda3/envs/py3b/lib/python3.6/multiprocessing/process.py", line 93, in run
    self._target(*self._args, **self._kwargs)
  File "/home/vision/lpx/AE23D/test_pycuda.py", line 22, in ffunc
    block=(400,1,1), grid=(1,1))
  File "/home/vision/anaconda3/envs/py3b/lib/python3.6/site-packages/pycuda/driver.py", line 382, in function_call
    func._set_block_shape(*block)
pycuda._driver.LogicError: cuFuncSetBlockShape failed: initialization error
{}

Process finished with exit code 0

我不确定 pycuda 是否可以在不同的进程中运行。我期待您的建议。

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

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

multiprocessing.set_start_method('spawn')