使用Python多处理和pyfft用GPU计算几个FFT：如何避免GPU内存泄漏？

Question

我尝试在Python中使用pycuda和pyfft包为多CPU和单GPU计算实现以下模式。在

我希望有几个流程（例如多处理.池（）），每个人都能使用GPU（使用NVIDIA CUDA）执行FFTs。在

但是，我有以下问题：

如果我运行的进程太多或每个进程运行的fft太多，整个脚本将保持不变，而不会终止（也不会计算所有到期的fft）。根据进一步的调查，我认为这是由于GPU上的内存限制（NVIDIA GeForce GT 750M目前为2048MB）。似乎多处理池无法获得控制权。 有什么办法可以避免这种情况吗？在

因为每个进程需要的内存少于2048MB，所以我希望有一个队列，每个进程都可以预订GPU的使用情况，当一个进程释放上下文时，队列中的下一个进程开始使用它。 这可行吗？在

或者，是否可以强制在给定时间只有一个进程使用GPU？
我分别尝试过这些解决方案，但它们都不起作用（或者我可能没有正确地实施它们）：

1. 使用proc同步流_流.同步（）
2. 清除上下文缓存，使用pycuda.tools.clear_上下文缓存（）
3. 更改计算模式，使用cuda.compute U模式= cuda.compute U模式.独家

注意：解决方案2。似乎释放了一些内存，但这会使计算速度变慢，并且不能解决问题：例如，增加要计算的fft的数量，脚本会显示相同的行为。在

这里是密码。为了从一个简单的任务开始，这里每个进程计算1个FFT（然后可以使用execute（）中的批处理选项在一行中执行更多FFT）。在

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

def main():
    # generates simple matrix, (e.g. image with a signal at the center)
    size = 4096
    center = size/2
    in_matrix = np.zeros((size, size), dtype='complex64')
    in_matrix[center:center+2, center:center+2] = 10.

    pool_size = 4  # integer up to multiprocessing.cpu_count()
    pool = multiprocessing.Pool(processes=pool_size)
    func = FuncWrapper(in_matrix, size)
    nffts = 16  # total number of ffts to be computed
    par = np.arange(nffts)

    results = pool.map(func, par)
    pool.close()
    pool.join()

    print results

这里是函数包装器：

^{pr2}$

现在，当nffts=16且pool_size=4时，脚本将正确终止并给出以下输出：

Func initialized with matrix size=4096
PoolWorker-1 free=1481.019392   tot=2147.024896
PoolWorker-2 free=1331.011584   tot=2147.024896
PoolWorker-3 free=1181.003776   tot=2147.024896
PoolWorker-4 free=1030.631424   tot=2147.024896
PoolWorker-1 free=881.074176    tot=2147.024896
PoolWorker-2 free=731.746304    tot=2147.024896
PoolWorker-3 free=582.418432    tot=2147.024896
PoolWorker-4 free=433.090560    tot=2147.024896
PoolWorker-1 free=582.754304    tot=2147.024896
PoolWorker-2 free=718.946304    tot=2147.024896
PoolWorker-3 free=881.254400    tot=2147.024896
PoolWorker-4 free=1030.684672   tot=2147.024896
PoolWorker-1 free=868.028416    tot=2147.024896
PoolWorker-2 free=731.713536    tot=2147.024896
PoolWorker-3 free=582.402048    tot=2147.024896
PoolWorker-4 free=433.090560    tot=2147.024896
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

但是当nffts=18且pool_size=4时，脚本不会终止并给出以下输出，仍然停留在最后一行：

Func initialized with matrix size=4096
PoolWorker-1 free=1416.392704   tot=2147.024896
PoolWorker-2 free=982.544384    tot=2147.024896
PoolWorker-1 free=1101.037568   tot=2147.024896
PoolWorker-2 free=682.991616    tot=2147.024896
PoolWorker-3 free=815.747072    tot=2147.024896
PoolWorker-4 free=396.918784    tot=2147.024896
PoolWorker-3 free=503.046144    tot=2147.024896
PoolWorker-4 free=397.144064    tot=2147.024896
PoolWorker-1 free=531.361792    tot=2147.024896
PoolWorker-1 free=397.246464    tot=2147.024896
PoolWorker-2 free=518.610944    tot=2147.024896
PoolWorker-2 free=397.021184    tot=2147.024896
PoolWorker-3 free=517.193728    tot=2147.024896
PoolWorker-4 free=397.021184    tot=2147.024896
PoolWorker-3 free=504.336384    tot=2147.024896
PoolWorker-4 free=149.123072    tot=2147.024896
PoolWorker-1 free=283.340800    tot=2147.024896

非常感谢你的帮助！在