使用 cupy 进行异步 GPU 内存传输

Question

是否可以使用 cupy（或 chainer）异步传输内存 from/to GPU？

我正在训练一个相对较小的网络，其中包含不适合 GPU 内存的非常大的数据。此数据应保存在 CPU 内存中，并按顺序提供给 GPU 进行小批量计算。

内存传输时间是该应用程序的主要瓶颈。我认为异步内存传输解决了这个问题，即在一个minibatch的计算过程中，另一个minibatch在后台传输到GPU。

我想知道 cupy.cuda.Stream class 是否可行，但我还不知道。我将不胜感激 comments/advice.

编辑：我以为下面的代码可以实现异步内存传输，但不是。

import numpy as np
import cupy as cp

a_cpu = np.ones((10000, 10000), dtype=np.float32)
b_cpu = np.ones((10000, 10000), dtype=np.float32)

a_stream = cp.cuda.Stream(non_blocking=True)
b_stream = cp.cuda.Stream(non_blocking=True)

a_gpu = cp.empty_like(a_cpu)
b_gpu = cp.empty_like(b_cpu)

a_gpu.set(a_cpu, stream=a_stream)
b_gpu.set(b_cpu, stream=b_stream)

# This should start before b_gpu.set() is finished.
a_gpu *= 2

nvvp 显示内存传输是按顺序进行的。

Answer 1

我通过深入研究 chainer source code 找到了一个解决方案。

在构造np.ndarray的时候，似乎要保持一个固定的内存缓冲区是很重要的一点。

def pinned_array(array):
    # first constructing pinned memory
    mem = cupy.cuda.alloc_pinned_memory(array.nbytes)
    src = numpy.frombuffer(
                mem, array.dtype, array.size).reshape(array.shape)
    src[...] = array
    return src

a_cpu = np.ones((10000, 10000), dtype=np.float32)
b_cpu = np.ones((10000, 10000), dtype=np.float32)
# np.ndarray with pinned memory
a_cpu = pinned_array(a_cpu)
b_cpu = pinned_array(b_cpu)

a_stream = cp.cuda.Stream(non_blocking=True)
b_stream = cp.cuda.Stream(non_blocking=True)

a_gpu = cp.empty_like(a_cpu)
b_gpu = cp.empty_like(b_cpu)

a_gpu.set(a_cpu, stream=a_stream)
b_gpu.set(b_cpu, stream=b_stream)

# wait until a_cpu is copied in a_gpu
a_stream.synchronize()
# This line runs parallel to b_gpu.set()
a_gpu *= 2

使用 cupy 进行异步 GPU 内存传输

Asynchronous GPU memory transfer with cupy

python

chainer

cupy