cupyによる非同期GPUメモリ転送

Question

cupy（またはchainer）でGPUとの間で非同期にメモリを転送できますか？

私は、GPUのメモリに収まらない非常に大きなデータを持つ比較的小さなネットワークを訓練しています。 このデータはCPUメモリに保持され、ミニバッチ計算のためにGPUに提供されます。

メモリ転送時間は、このアプリケーションの支配的なボトルネックです。 私は、非同期メモリ転送がこの問題を解決したと思います。すなわち、あるミニバッチの計算中に 、もう1つのミニバッチがバックグラウンドで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 source codeにダイビングすると1つの解決策が見つかりました。

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