如何使用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)

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 是否可以在不同的進程中運行。 我期待着您的建議。

Answer 1

不應在 fork 之前初始化 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')工作完美。

Answer 2

幸運的是，我解決了這個問題。

在main函數中添加一行代碼：

multiprocessing.set_start_method('spawn')