NumPyからCuPyへ：高速化の一例

1 はじめに
2 環境
3 インストール
4 CuPyについて
5 CuPyの苦手な処理
6 まとめ

はじめに

Pythonで数値計算を行う際によく使われるライブラリがNumPyです。しかし、GPUを活用することで計算速度を向上させたい場面も多いでしょう。そこでCuPyが登場します。この記事では、NumPyのコードをCuPyに置き換える一例とその性能比較について解説します。一方でCuPyの苦手な処理についても触れますので、使いどころを間違えないようにしましょう。

環境

$ pip install numpy cupy
Installing collected packages: fastrlock, numpy, cupy
Successfully installed cupy-12.2.0 fastrlock-0.8.2 numpy-1.24.4
(cupy)
$ pip list
Package       Version
------------- -------
cupy          12.2.0
fastrlock     0.8.2
numpy         1.24.4
pip           23.2.1
pkg_resources 0.0.0
setuptools    44.0.0

$ python -V
Python 3.8.10
$ uname -a
Linux user 5.15.0-83-generic #92~20.04.1-Ubuntu SMP Mon Aug 21 14:00:49 UTC 2023 x86_64 x86_64 x86_64 GNU/Linux

Ubuntu 20.04

インストール

バイナリパッケージ（ホイール形式）は、PyPI上でLinuxおよびWindows向けに利用可能。

Platform	Architecture	Command
CUDA 10.2	x86_64 / aarch64	`pip install cupy-cuda102`
CUDA 11.0	x86_64	`pip install cupy-cuda110`
CUDA 11.1	x86_64	`pip install cupy-cuda111`
CUDA 11.2 ~ 11.8	x86_64 / aarch64	`pip install cupy-cuda11x`
CUDA 12.x	x86_64 / aarch64	`pip install cupy-cuda12x`
ROCm 4.3 (experimental)	x86_64	`pip install cupy-rocm-4-3`
ROCm 5.0 (experimental)	x86_64	`pip install cupy-rocm-5-0`

上記を参考に、pipでインストールします。

CuPyについて

CuPyはPythonプログラミング言語でGPUによる高速計算をサポートするオープンソースライブラリです。多次元配列、疎行列、およびそれらの上で実装された多様な数値アルゴリズムをサポートしています。CuPyはNumPy同じAPIセットを共有しており、NumPy/SciPyのコードをGPUで実行するためのドロップイン置換として機能します。CuPyはNVIDIAのCUDA GPUプラットフォームと、v9.0からはAMDのROCm GPUプラットフォームもサポートしています。

https://github.com/cupy/cupy

NumPyのサンプルコード

まずは、行列の積を計算するシンプルなNumPyのコードを見てみましょう。

import time

import numpy as np

# タイム計測開始
start_time = time.time()

# 行列生成
a = np.random.rand(5000, 5000)
b = np.random.rand(5000, 5000)

# 行列の積
result = np.dot(a, b)

# タイム計測終了
end_time = time.time()

print(f"NumPy Time: {end_time - start_time} seconds")

CuPyによる高速化

次に、上記のコードをCuPyに置き換えてみます。

import time

import cupy as cp

# タイム計測開始
start_time = time.time()

# 行列生成
a = cp.random.rand(5000, 5000)
b = cp.random.rand(5000, 5000)

# 行列の積
result = cp.dot(a, b)

# タイム計測終了
end_time = time.time()

print(f"CuPy Time: {end_time - start_time} seconds")

性能比較

両者のコードを実行した結果、CuPyの方が計算速度が大幅に向上したことが確認できました。

NumPy Time: 3.5 seconds
CuPy Time: 1.0 seconds
初回オーバーヘッド: CuPy（または他のGPUライブラリ）をはじめて使用する際には、GPUの初期化などに時間がかかる場合があります。このオーバーヘッドは一度だけ発生することが多いです。

CuPyの苦手な処理

CuPyを使用する際には、CPUからGPUへのデータ転送が必要なケースがあります。このデータ転送は、大量のデータを扱う場合には、パフォーマンスに悪影響となります。以下に、この「苦手な処理」を模倣する簡単なPythonコードを示します。

NumPyのサンプルコード

import time

import numpy as np

# NumPyで大量のデータを生成
n_elements = 10**7  # 要素数
numpy_array = np.random.rand(n_elements)

# 何らかのCPU処理（例：要素ごとの平方根の計算）
start_time_calc = time.time()  # 計算タイム計測開始
result_numpy = np.sqrt(numpy_array)
end_time_calc = time.time()  # 計算タイム計測終了

print(f"NumPy Calculation time: {end_time_calc - start_time_calc} seconds")  # 計算にかかった時間を表示

CuPyのサンプルコード

import time

import numpy as np

import cupy as cp

# NumPyで大量のデータを生成
n_elements = 10**7  # 要素数
numpy_array = np.random.rand(n_elements)

# CPUからGPUへのデータ転送（これが苦手な処理）
start_time_transfer = time.time()  # 転送タイム計測開始
cupy_array = cp.asarray(numpy_array)  # CPUからGPUへ転送
end_time_transfer = time.time()  # 転送タイム計測終了

print(f"CuPy Data transfer time: {end_time_transfer - start_time_transfer} seconds")  # 転送にかかった時間を表示

# 何らかのGPU処理（例：要素ごとの平方根の計算）
start_time_calc = time.time()  # 計算タイム計測開始
result_cupy = cp.sqrt(cupy_array)
end_time_calc = time.time()  # 計算タイム計測終了

print(f"CuPy Calculation time: {end_time_calc - start_time_calc} seconds")  # 計算にかかった時間を表示

性能比較

NumPy Calculation time: 0.02361893653869629 seconds
CuPy Data transfer time: 0.2935044765472412 seconds
CuPy Calculation time: 0.34106016159057617 seconds

約17倍の差があります。このように、データ転送が必要なケースでは、CuPyのパフォーマンスが悪化することがあります。
使いどころを間違えないようにしましょう。