pyCUDAで画像処理に挑戦してみた（超入門編）

どうもこんにちは。

ちゃんとブログを更新しようと思っている今日この頃です。

前回の記事で、GPU環境をいい感じに整えることができるnVidia-dockerについてご紹介しました。

mhr380.hatenablog.com

せっかく作って壊せるGPU環境が手に入ったので（？） 今回は、CUDAを使ったGPUプログラミングをやってみたいと思います。 ところで、CUDAで遊んでみたいけど、C++から叩くのは面倒なあ… という方は結構いらっしゃるのでは無いでしょうか。 もちろん、僕もそのタイプです。

そんな僕にピッタリのツールとしてpyCudaがあります。

Welcome to PyCUDA’s documentation! — PyCUDA 2017.1.1 documentation

PythonからCUDAのカーネルを呼ぶことができるライブラリです。しかし、情報があまりないんですよね…

今回は練習がてら、numpy array形式の画像をGPUで処理するようなコードを書いてみたいと思います。

環境

以下の環境で動作確認をしています。

• Ubuntu16.04 LTS

• GeForce GTX1070

• nvidia-docker

準備

nVidia-dockerのコンテナとして環境を用意することにします。 ベースとして、以下のイメージを利用しました。

https://hub.docker.com/r/nightseas/pycuda/

ただ、上記イメージにはPillowが入っていなかったので足してあげます。 私の環境では、下記Dockerfileをビルドすると正常に動作しました。 ただし、ベースとしたイメージはPython2.7.12だったので、3系を使いたい方は別途環境を構築する必要がありそうです。 私もそのうちやってみます。

FROM nightseas/pycuda:latest

RUN apt-get update && \
    apt-get install -y python-tk

RUN pip install pillow

やってみること

最初は基礎的なところから始めましょう。 8bitグレースケール画像の輝度を反転するだけのコードを書いてみたいと思います。

実装

# coding: utf-8

import pycuda.autoinit
import pycuda.driver as cuda
from pycuda.compiler import SourceModule

import numpy as np
from PIL import Image


#
# 以下CUDAカーネル
#

mod = SourceModule("""
__global__ void negative(unsigned char *image, int height, int width)
{
    int pos_y = threadIdx.y + blockDim.y * blockIdx.y; // 画像のy方向の位置を取得。
    int pos_x = threadIdx.x + blockDim.x * blockIdx.x; // 画像のx方向の位置を取得。

    int idx = pos_y * height + pos_x; // 内部では1次元配列化しているので、インデックスはこのようになる

    image[idx] = 255 - image[idx];  //画素値を反転
}
""")


#
# 画像の読みこみ
#
pil_image = Image.open("./lena_gray.png")   # 画像の読み込み
image = np.array(pil_image, np.uint8)       # PIL Image型からnumpy arrayに変換。

height, width = image.shape[:2] # 今回は512x512pxの画像を利用

#
# CUDA周りの準備
#
cuda_kernel = mod.get_function("negative")  # 上で定義したカーネルを呼び出す

block   = (512, 1, 1)                       # Block size (後述) 
grid    = (512, 1, 1)                      # Grid size (後述)


#
# カーネルを実行
#
cuda_kernel(cuda.InOut(image),              # imageを参照渡しする。
            np.int32(height),               # int型定数はnumpyで明示的に型を定義する
            np.int32(width), 
            block=block, grid=grid)         # BlockとGridも引数として与える


#
# 処理後の画像を保存
#
pil_output_image = Image.fromarray(image)
pil_output_image.save("out.png")

CUDAにおける画素インデックスは、下記スライドのp59を見るとわかりやすいと思います。

CUDAプログラミング入門 from NVIDIA Japan

www.slideshare.net

こんな感じで初歩のの初歩ですがCUDAを使った画像処理を行なうことができます。 次回（あるのか？）は、もうちょっと凝ったことにチャレンジしたいと思います。

2017/08/29追記 qiita.com Gridsize, Blocksizeは上記の記事を参考に修正しました。