numpyを使ったプログラムをGPUで動かす方法

通常のプログラムは CPU で計算していますが,機械学習やディープラーニングでは GPU を用いればより早く計算することができる(場合もある)ことはよく知られています.

ハードルが高いように思えますが,実は Numpy で実装されているような機能であれば CuPy という強力なパッケージが GPU がサポートしてくれます.

ディープラーニングであれば TensorFlow, PyTorch, Keras などのパッケージを用いることで GPU にも対応するようなプログラムを書くことができますが,単純な行列計算等のためにこれらのパッケージをわざわざ使うのも,変数の融通などを考えると微妙な気がします.

CuPy とは

Chainer を開発していたことでも有名な PFN が提供している演算用パッケージです.

NVIDIA GPU では CUDA ソフトウェアを用いて GPGPU(GPU による汎用計算)が可能ですが,CUDA でプログラムを動かすためには,通常と異なる特殊な書き方をしなければなりません.

CuPy はこの CUDA と Python の橋渡しをしてくれます.Numpy の多くの関数をサポートしているので,それとほとんど変わらない書き方で GPGPU が実現できます.

必須環境

  • CUDA toolkit: 8.0 ~
  • Python: 3.5.1 ~
  • Numpy: 1.9 ~

推奨 OS は Ubuntu 16.04/18.04, CentOS 7 で,Windows でも動作可能なようです.最近の macOS 搭載 PC は Radeon GPU なので CUDA が対応していません.

Experimental で ROCm (Radeon Open Compute) 対応もありますが,推奨 OS は Ubuntu 16.04/18.04 となっています.

macOS は無理っぽいです.NVIDIA は CUDA 最新バージョンで macOS のサポートを切ったとかなんとか.macOS は PlaidML+各種ライブラリか OpenCL なら GPGPU いける...?

詳細は Requirements で確認してください.

インストール

CUDA が既にインストールされているものとして.

terminal

$ pip install cupy

基本的な書き方

Numpy に準拠

CuPy は Numpy で定義された関数のほとんどをサポートしていて,同じ関数名で同じ機能を使うことができます.素晴らしすぎる!!

import cupy as cp
import numpy as np

# Compute on CPU
W_cpu = np.random.randint(1, 10, (3, 4))
x_cpu = np.array([1, 3, 2, 4])[:, np.newaxis]
y_cpu = np.dot(W_cpu, x_cpu)

# Compute on GPU
W_gpu = cp.random.randint(1, 10, (3, 4))
x_gpu = cp.array([1, 3, 2, 4])
[:, np.newaxis]
y_gpu = cp.dot(W_gpu, x_gpu)

注意点

一つだけ気をつけなければならないのは,CPU での計算と GPU での計算では参照するメモリが異なることです.

import cupy as cp
import numpy as np

W_cpu = np.random.randint(1, 10, (3, 4))
x_cpu = np.array([1, 3, 2, 4])[:, np.newaxis]
# Can't compute
y_gpu = cp.dot(W_cpu, x_cpu)

data_gpu = cp.loadtxt('sample.csv', delimiter=',')
# Can't execute
np.savetxt('sample2.csv', data_gpu, delimiter=',')

このように CPU が参照するメモリに載ったデータや変数をそのまま GPU で計算することはできません.これは計算だけでなく,ファイルに保存するときなども同様です.

CPU と GPU で相互にデータや変数をやり取りするには,cupy.asarray()cupy.asnumpy() を使います..

import cupy as cp
import numpy as np

W_cpu = np.random.randint(1, 10, (3, 4))
x_cpu = np.array([1, 3, 2, 4])[:, np.newaxis]
# Move array to a device
W_gpu = cp.asarray(W_cpu)
W_gpu = cp.asarray(x_cpu)
y_gpu = cp.dot(W_gpu, x_gpu)

data_gpu = cp.loadtxt('sample.csv', delimiter=',')
# Move array from a device to the host
data_cpu = cp.asnumpy(data_gpu)
np.savetxt('sample2.csv', data_cpu, delimiter=',')

GPU メモリが足りないときの応急処置

GPU を使おうと思ったとき,一枚のメモリはだいたい 16GB,多くても 32GB のことが多いです.(メモリサイズ ∝ 価格)

そのため,データが載り切らない場合は行列を分割したり,複数の GPU にタスクを明示的に割り振ったりする必要があります.

それが適わない場合,応急処置として Unified Memory を使う方法もあります.

import cupy as cp

# Declare at first
pool = cp.cuda.MemoryPool(cp.cuda.malloc_managed)
cp.cuda.set_allocator(pool.malloc)

GPU と CPU の間でうまくデータを融通することでメモリの不足分を補うことができますが,データ転送が増えるのでパフォーマンスは低下します.

(Google Colaboratory では使用不可)