1

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

我想要几个进程（例如使用 multiprocessing.Pool() 启动），每个进程都能够使用 GPU（使用 NVIDIA CUDA）执行 FFT。

但是，我有以下问题：

如果我运行太多进程或每个进程有太多 FFT，则整个脚本将保持暂停状态而不会终止（并且不会计算所有到期的 FFT）。通过进一步的调查，我认为这是由于 GPU 的内存限制（目前 NVIDIA GeForce GT 750M 上的内存限制为 2048MB）。似乎多处理池无法重新获得控制权。有没有办法避免这种情况？

由于每个进程需要少于 2048 MB，我希望有一个类似队列的东西，每个进程都可以预订GPU 的使用，当一个进程释放上下文时，队列中的下一个进程开始使用它。这是可行的吗？

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

  1. 使用 proc_stream.synchronize() 同步流
  2. 使用 pycuda.tools.clear_context_caches() 清除上下文缓存
  3. 更改计算模式，使用 cuda.compute_mode = cuda.compute_mode.EXCLUSIVE

注意：解决方案 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

这里是函数包装器：

class FuncWrapper(object):
    def __init__(self, matrix, size):
        self.in_matrix = matrix
        self.size = size
        print("Func initialized with matrix size=%i" % size)

    def __call__(self, par):
        proc_id = multiprocessing.current_process().name

        # take control over the GPU
        cuda.init()
        context = make_default_context()
        device = context.get_device()
        proc_stream = cuda.Stream()

        # move data to GPU
        # multiplication self.in_matrix*par is just to have each process computing
        # different matrices
        in_map_gpu = gpuarray.to_gpu(self.in_matrix*par)

        # create Plan, execute FFT and get back the result from GPU
        plan = Plan((self.size, self.size), dtype=np.complex64,
                    fast_math=False, normalize=False, wait_for_finish=True,
                    stream=proc_stream)
        plan.execute(in_map_gpu, wait_for_finish=True)
        result = in_map_gpu.get()

        # free memory on GPU
        del in_map_gpu

        mem = np.array(cuda.mem_get_info())/1.e6
        print("%s free=%f\ttot=%f" % (proc_id, mem[0], mem[1]))

        # release context
        context.pop()

        return par

现在，使用 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

非常感谢您的帮助！

4

0 回答 0