如何在多處理中使用 PyCuda？

Question

問題不是在進程中使用 CUDA，而是如果父進程已經有一個初始化的 CUDA，那么進程中就不能初始化它。 如果您只是在父級中進行初始化並將其扔到進程中，但它不起作用。 以下是流程中的一個示例：

from pyfft.cuda import Plan
import numpy
import pycuda.driver as cuda
from pycuda.tools import make_default_context
import pycuda.gpuarray as gpuarray

from multiprocessing import Process

def do_this_fft(data):
    cuda.init()
    context = make_default_context()
    stream = cuda.Stream()
    plan = Plan((16, 16), stream=stream)

    gpu_data = gpuarray.to_gpu(data)
    plan.execute(gpu_data)
    result = gpu_data.get()
    del gpu_data
    print(result)
    context.pop()
    del plan,stream,context

data = numpy.ones((16, 16), dtype=numpy.complex64)
process_list=[Process(target=do_this_fft,args=(data,)).start() for i in range(2)]

但是如果我們寫這樣的代碼，那么什么都行不通了

from pyfft.cuda import Plan
import numpy
import pycuda.driver as cuda
from pycuda.tools import make_default_context
import pycuda.gpuarray as gpuarray

from multiprocessing import Process

def start_cuda_and_fft(data):
    cuda.init()
    context = make_default_context()
    stream = cuda.Stream()
    plan = Plan((16,16), stream=stream)
    context.pop()
    del plan,stream,context
    process_list=[Process(target=do_this_fft,args=(data,plan)).start() for i in range(2)] 

def do_this_fft(data):
    cuda.init() #<---------------Error
    '''
    File "test.py", line 35, in do_this_fft
        cuda.init()
    LogicError: cuInit failed: initialization error

    '''
    context = make_default_context()
    stream = cuda.Stream()
    plan = Plan((16, 16), stream=stream)

    gpu_data = gpuarray.to_gpu(data)
    plan.execute(gpu_data)
    result = gpu_data.get()
    context.pop()
    del plan,stream,context
    print(result)

data = numpy.ones((16, 16), dtype=numpy.complex64)
process_list=[Process(target=do_this_fft,args=(data,)).start() for i in range(2)]

錯誤如下：

文件“test.py”，第 35 行，在 do_this_fft cuda.init() LogicError: cuInit failed: 初始化錯誤

我已經嘗試做上下文、計划和推進過程，所有這些都沒有任何后果。

Answer 1

我們不能在 fork() 之前初始化 CUDA 所以我們需要產生一個新進程

Answer 2

我有一個類似的問題，並通過在啟動子進程之前在主進程上添加一行代碼來解決它：

 multiprocessing.set_start_method('spawn')

來源：[ https://stackoverflow.com/a/55812288/8664574]