1.8 CUDAによる学習の高速化

「ゼロから作るDeep Learning」では第2巻でCuPyを導入しています。Ivoryライブラリでも同様にCUDAを使った学習ができる設計となっています。

なお、Windows環境では、CuPyのインストール時にエラーが発生する場合があります。下記のようにすると回避できます。(使用するCUDAのバージョンに合わせてパッケージ名を変更する必要があります。)

conda install fastrlock
pip install cupy-cuda100

実際に使ってみます。

import cupy as cp

a = cp.arange(10)
print(type(a))
print(a.dtype)
print(a.device)

[1] 2019-06-12 20:00:40 (1.12s) python3 (1.12s)

<class 'cupy.core.core.ndarray'>
int32
<CUDA Device 0>

Ivoryライブラリでは、ivory.common.contextモジュールを使ってCPUとCUDAを切り替えます。

from ivory.common.context import np

[2] 2019-06-12 20:00:42 (35.6ms) python3 (1.16s)

コンテキスを確認します。

np.context

[3] 2019-06-12 20:00:42 (7.97ms) python3 (1.17s)

'cpu'

NumPyと同じように使います。

a = np.arange(10, dtype="f")
print(type(a))
print(a.dtype)

[4] 2019-06-12 20:00:42 (10.0ms) python3 (1.18s)

<class 'numpy.ndarray'>
float32

コンテキストを変更して、同じコードを試してみましょう。

np.context = "cuda"
a = np.arange(10, dtype="f")
print(type(a))
print(a.dtype)

[5] 2019-06-12 20:00:42 (13.0ms) python3 (1.19s)

<class 'cupy.core.core.ndarray'>
float32

今度はCuPyのアレイが得られました。このように同じコードで切り替えが可能です。参考までに、contextモジュールをソースを記載しておきます。

File 1.1 ivory.common.contextモジュールファイル

import numpy

numpy.scatter_add = numpy.add.at

try:
    import cupy
except ImportError:
    cupy = None
else:
    pass
    # cupy.cuda.set_allocator(cupy.cuda.MemoryPool().malloc)


class Context:
    def __init__(self):
        self.context = "cpu"

    def __getattr__(self, name: str):
        if self.context == "cpu":
            return getattr(numpy, name)
        else:
            return getattr(cupy, name)


np = Context()