CuPy を使ってみる

2017年11月3日

はじめに

CuPy を使ってみる。

環境

• Ubuntu 16.04 LTS
• Anaconda 5.0.1 (Python 3.6)
• CUDA 9.0

CuPy のインストール

pip でインストールする。

$ pip install cupy

CUDA のセットアップが済んでいないと失敗するはず。

CuPy を使ってみる

CuPy は NumPy の代わりとして使えるようである。n x n の行列積で NumPy と CuPy で速度を比較してみる。

cupy_test.py

import sys
import time
import numpy as np
import cupy as cp

argv = sys.argv
argc = len(argv)

if argc < 2:
    print('usage: cupy_test.py <n>')
    exit(0)

n = int(sys.argv[1])

print('n = %d' % n)

x = np
start = time.time()
a = x.random.rand(n, n)
b = x.random.rand(n, n)
c = a*b
end = time.time()
numpy_time = end - start
print('NumPy: %f s' % numpy_time)

x = cp
start = time.time()
a = x.random.rand(n, n)
b = x.random.rand(n, n)
c = a*b
end = time.time()
cuda_time = end - start
print('CuPy: %f s' % cuda_time)
print('x %f' % (numpy_time/cuda_time))

つぎのように n を指定して実行する。

$ python cupy_test.py 1000

n = 1000 と n = 10000 について、3 回ずつ計算して計算時間の平均を取った。

n が小さい場合、CuPy は CUDA の準備に時間がかかりすぎてかえって遅くなるようだが、n が十分に大きければ NumPy の 5 倍程度の速度が出るようである。