與cupy異步GPU內存傳輸

Question

是否可以通過cupy （或chainer ）從GPU異步傳輸內存？

我正在訓練一個相對較小的網絡，其數據非常大，不適合GPU內存。 此數據應保存在CPU內存中，並提供給GPU以按順序進行小批量計算。

內存傳輸時間是此應用程序的主要瓶頸。 我認為異步內存傳輸解決了這個問題，即在計算一個小批量時，另一個小批量在后台轉移到GPU。

我想知道它可以用cupy.cuda.Stream類，但我還不知道。 我將不勝感激任何意見/建議。

編輯：我認為以下代碼進行異步內存傳輸，但不是。

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源代碼找到了一個解決方案。

在構造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