【GPUSOROBAN】PyCudaの使用例:CPUとGPUにおけるPyTorchの処理時間比較

Contents

• 1 はじめに
• 2 CPUとGPUにおけるPyTorchの処理時間
  • 2.1 PyTorchでGPU情報を確認
• 3 計算の流れ
• 4 2次元配列の四則演算

はじめに

GPUとはGraphics Processing Unitの略で、描画処理用のCPUのことを指しています。元々はGraphic用途だったのを汎用的な計算に拡大することをGPGPUと呼ばれています。General Purpose Graphics Processing Unit(GPGPU)はGPUを画像処理ではなく、他の用途（ex. 暗号解読、音声処理）に使用するときに使用されています。今回Nvidia社が作ったCUDAを機械学習の裏でどのように実行されているかを紹介いたします。CUDAとはGPU並列を扱うためのプログラミング言語です。C言語を拡張したものはPyCUDAです(CUDAのPython Wrapper API)。Kaggleや研究などでは、CUDAをPythonから扱えば、メモリ管理が楽になります。

実行環境：conda_torch17_py36, PyCuda version : 2021.1

CPUとGPUにおけるPyTorchの処理時間

PyTorchでGPU情報を確認

PyTorchでGPUの情報を取得する関数はtorch.cuda以下に用意されています。GPUが使用可能かを確認するtorch.cuda.is_available()を使って、自分の環境でGPUが使えているかどうかを確認します。

In [1]:

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

print("using", device, "device")

using cuda device

GPUが使えることを確認したら、簡単なCPUとGPUの処理時間を見比べましょう。GPUで計算を行う際に、GPUにデータをPyTorchの.to(device)を使って転送する必要があります。 また、GPUの計算は裏で行われるため、計算の処理が終わるまで、torch.cuda.synchronize()でプログラムの待機が必要です。初めてPyCUDAを実行するとき、GPUの初期化などの処理が必要になるため、処理時間が遅くなるけど、２回目以降にGPUの計算を行うと、早くなることがわかりました。そのため、今回３回ぐらいの計算を行って、その平均値を取って、CPUの処理時間と比べます。

In [2]:

import time

matrix_size = 30*512

x = torch.randn(matrix_size, matrix_size)
y = torch.randn(matrix_size, matrix_size)

print("************ CPU SPEED ***************")
start = time.time()
result = torch.matmul(x,y)
print(time.time() - start)
print("verify device:", result.device)

x_gpu = x.to(device)
y_gpu = y.to(device)
torch.cuda.synchronize()

for i in range(3):
    print("************ GPU SPEED ***************")
    start = time.time()
    result_gpu = torch.matmul(x_gpu,y_gpu)
    torch.cuda.synchronize()
    print(time.time() - start)
    print("verify device:", result_gpu.device)

CPU SPEED 
296.8993308544159
verify device: cpu
GPU SPEED 
0.5544977188110352
verify device: cuda:0
GPU SPEED 
0.4655277729034424
verify device: cuda:0
* GPU SPEED 
0.45668745040893555
verify device: cuda:0

PyTorchでGPUによる計算が早いことがわかりました。しかし、 CUDAでの計算流れがどのように実行されますでしょうか？

計算の流れ

In [3]:

## 必要なライブラリをImportする

import numpy as np
import math
import pycuda.gpuarray as gpuarray
from pycuda.elementwise import ElementwiseKernel
from pycuda.compiler import SourceModule

import pycuda.autoinit

PyCUDAのコードを書きます。今回ElementwiseKernelを使って、CUDAカーネルを記述してPythonから実行します。PythonのForループになっていた遅い部分をGPUで実行したい時によく用いられます。他のユースケースは以下となります。

• 配列（行列・テンソル）の要素ごとに同じような処理したい
• バッチで入力した画像全部に同じような処理したい

ElementwiseKernelを使えば、CUDAカーネルを簡易に書くことができます。実行時にカーネルの実行内容はコンパイルされます。詳細はCuPyの参考資料をご覧ください。

In [4]:

plus_one_kernel = ElementwiseKernel(
"int *x, int *y",
"y[i] = x[i] + 1",
"plus_one")

np.arangeは、連番や等差数列を生成する関数です。np.arangeを使って、配列xとyを作成します。

In [5]:

num_components = 10
x = np.arange(num_components, dtype=np.int32)

CPUからGPUに配列を転送します。

In [6]:

x_gpu = gpuarray.to_gpu(x)
y_gpu = gpuarray.zeros(num_components, dtype=np.int32)

In [7]:

plus_one_kernel(y_gpu, x_gpu)

結果を取得します。

In [8]:

y_gpu.get()

Out[8]:

array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)

In [9]:

x

Out[9]:

array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)

2次元配列の四則演算

SourceModelはCUDA Cカーネルをコンパイルする関数SourceModelにより一般的なCUDAカーネルを実行することができます。詳細はCuPyの参考資料をご覧ください。

In [10]:

module = SourceModule("""
__global__ void add_two_array_2d(int nx, int ny, float *res, float *arr1, float *arr2){
  int x = threadIdx.x + blockDim.x * blockIdx.x;
  int y = threadIdx.y + blockDim.y * blockIdx.y;
  int ij = nx * y + x;
  if (x < nx && y < ny){
      res[ij] = arr1[ij] + arr2[ij]; 
  }
}
""")

In [11]:

add_two_array = module.get_function("add_two_array_2d")

In [12]:

num_x, num_y = np.int32(5), np.int32(2)
num_components = num_x * num_y

In [13]:

arr1 = np.arange(num_components, dtype=np.float32).reshape(num_y, num_x)

In [14]:

np.random.seed(123)
arr2 = 10 * np.random.rand(num_y, num_x)
arr2 = arr2.astype(np.float32)

In [15]:

res = np.zeros([num_y, num_x], dtype=np.float32)

In [16]:

arr1_gpu = gpuarray.to_gpu(arr1)
arr2_gpu = gpuarray.to_gpu(arr2)
res_gpu = gpuarray.to_gpu(res)

In [17]:

threads_per_block = (16, 16, 1)
block_x = math.ceil(num_x / threads_per_block[0])
block_y = math.ceil(num_y / threads_per_block[1])
blocks_per_grid = (block_x, block_y, 1)

In [18]:

add_two_array(num_x, num_y, res_gpu, arr1_gpu, arr2_gpu, block=threads_per_block, grid=blocks_per_grid)

結果を取得します。

In [19]:

res_gpu.get()

Out[19]:

array
( [[ 6.9646916, 3.8613935, 4.2685146, 8.513147 , 11.19469 ],
[ 9.231065 , 15.807642 , 13.848297 , 12.809319 , 12.921175 ]],
dtype=float32)In [20]:

arr1 + arr2

Out[20]:

array
([[ 6.9646916, 3.8613935, 4.2685146, 8.513147 , 11.19469 ],
[ 9.231065 , 15.807642 , 13.848297 , 12.809319 , 12.921175 ]],
dtype=float32)