Python + Numba CUDA で GPU プログラミング入門以前 #Python

Python + Numba CUDA で GPU プログラミングしてみるシリーズ

その０： Python + Numba CUDA で GPU プログラミング入門以前 [← イマココ]
その１： Python + Numba CUDA で画像処理
その２： Python + Numba CUDA でライフゲーム

GPU マシンを買った¹ものの「TensorFlow を速く実行してくれるやつ」としてしか扱っておらず、もう少し使いこなせるようにならないと勿体ないので CUDA プログラミングを始めることにした。

今回は Python から CUDA を扱うためのツールとして Numba を使う。
Python から CUDA を扱うツールとしては他に PyCUDA や CuPy などがあるが、
NVIDIA の公式ページでは Numba が紹介されている² ので最初に触るものとしてはこれが良いのかと思った。

はじめに

この記事では「CUDA のプログラミングモデル」や「Numba の使い方」については触れない。ひとまず「Python コードを GPU 上で実行できるようになる」ことを目指す。

また CUDA プログラミングに Docker は必要ではないが、こういった実験的な遊びは環境を作って壊すことが容易にできたほうが良いと思っているので、今回は Docker を使っていく。

Ubuntu 16.04 に NVIDIA ドライバをインストール

割愛

Docker と NVIDIA Docker をインストール

割愛

Numba + CUDA が動く Docker イメージを作成

CUDA のイメージに Miniconda を入れて Numba と CUDA Toolkit をインストールする。³

Dockerfile

FROM nvidia/cuda:8.0-cudnn6-runtime

RUN apt-get update \
    && apt-get install -y wget bzip2 \
    && apt-get clean

RUN cd /tmp \
    && wget --quiet https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh \
    && /bin/bash miniconda.sh -b -p /opt/conda \
    && rm miniconda.sh

ENV PATH /opt/conda/bin:$PATH

RUN conda install -y -q cudatoolkit numba

Docker イメージをビルドして numba-cuda という名前を付ける。

$ docker build -t numba-cuda .

動かしてみる

GPU を使うので Docker コンテナを起動する際に --runtime=nvidia が必要。

$ docker run --runtime=nvidia --rm -it numba-cuda

ログインしたコンテナ内で以下の Python スクリプトを実行してみる。
以下は「配列のすべての要素に 1 を加える」という処理を GPU 上で行うスクリプト。

from numba import cuda
import numpy as np

@cuda.jit
def increment_by_one(arr):
    tx = cuda.threadIdx.x
    ty = cuda.blockIdx.x
    bw = cuda.blockDim.x
    pos = tx + ty * bw
    arr[pos] += 1

arr = np.arange(5)
print('Input: {}'.format(arr))

threadsperblock = 32
blockspergrid = 1
increment_by_one[blockspergrid, threadsperblock](arr)
print('Output: {}'.format(arr))

実行すると以下のような結果になる。

Input: [0 1 2 3 4]
Output: [1 2 3 4 5]

期待どおりに動いた。

これから

ひとまず簡単なスクリプトを動かせるようになったので、これから GPU や CUDA について理解を深めながらもっと複雑なことをやっていこうと思う。

今回はここまで。

http://blog.hotolab.net/entry/zbox ↩
GPU Accelerated Computing with Python | NVIDIA Developer ↩
conda を使わずとも pip で Numba をインストールできるが、CUDA を利用したコードは色々と試行錯誤しても動かせなかったため Miniconda を使って CUDA Toolkit を一緒に入れるほかなかった。そもそも CUDA のベースイメージに CUDA Toolkit が含まれているはずなのに、あとから conda でも入れないと動かないのは何故なのかが謎。 ↩