CUDAとプロセス並列化で超高速画像処理 with Julia

この記事でやりたいこと

NVIDIA GPUで大量の画像処理(20万枚の画像をフーリエ変換するとか)をしたいとき，CUDA C/C++ で書くと Glob を使いたくなるし，Python だと CuPy の遅さが気になるしで一長一短です．Juliaでいいとこ取りしましょう．さらに Distributed.jl の pmap() を使ってプロセス並列化することでファイルアクセスのオーバーヘッドを軽減でき，さらに高速化できます．

実装例

fftOnGPU

# マスタープロセスでだけ使うパッケージのロード
using Glob
using Distributed
# ワーカープロセスを起動．　今回は10並列．
addprocs(10)
# ワーカープロセスで行う処理は @everywhere をつける．
@everywhere using CUDA
@everywhere using Images
@everywhere using CUDA.CUFFT
@everywhere using FileIO

# 出力したい画像サイズの8bit BMPヘッダを予めpathに用意した画像から読み込む
function getBMPHeader(path)
    io = open(path,"r")
    header = read(io,1078)
    close(io)
    return header
end

# 8bit BMPを保存するための関数． あらかじめ用意したヘッダとデータで無理やり作るよ．
@everywhere function save8bitBMP(path,header,data,dataLen)
    io = open(path,"w")
    write(io,header)
    writeData = UInt8.(floor.(data*255))
    for i in 1:dataLen
        write(io,writeData[i,:])
    end
    close(io)
end

# 画像をGrayscaleで読み込んでFloat32に．
@everywhere function loadImg(path)
    out = Float32.(channelview(Gray.(load(path))))
end

# path の画像をGPUでFFTしてスペクトル分布を保存する関数
@everywhere function fftOnGPU(path, dataLen, header)
    # 処理するファイルのパスを表示
    println(path)
    # 白黒画像ファイルの読み込み．cu()でGPUに配列をコピー．
    img = cu(loadImg(path))
    # 二次元フーリエ変換と fftshift (画像中央が低周波数域になるように変換する処理)
    d_fft = CUDA.CUFFT.fftshift(CUDA.CUFFT.fft(img))
    # 諸々の処理．あとd_fftをComplexF32からFloat32に変換．
    d_fft = Float32.(log.(abs.(d_fft).+1))
    # 値域を0~1に
    d_fft = d_fft ./maximum(d_fft)
    # Array() でGPUからホストへ配列をコピー．
    hostout = Array(d_fft)
    # 保存用パスの準備
    newPath = replace(path,"hoge" => "hogefft")
    println(newPath)
    # 8bit BMPとして保存． Images.jlのsaveだと24bitBMPになっちゃう．
    save8bitbmp(newPath,header,hostout,dataLen)
end

# ./hoge ディレクトリのbmpファイルのパスをリストで取得
pathList = glob("./hoge/*.bmp")
# 画像一辺の長さ
dataLen = 1024

# sample.bmp は出力と同型の画像．ヘッダと流用します．
header = getBMPHeader("./sample.bmp")

# 10個のワーカープロセスに処理を分散．
pmap(x -> fftOnGPU(x, dataLen, header),pathList)

補足

• fft()やfftshift()だけCUDA.CUFFTをつけているのは，Images.jlのFFTWと競合するからです．
• Images.jl のsave()ではモノクロ8bitBMPの保存はできません．なので出力と同じサイズの8bit BMPを事前に(ImageJとかで)作成して，ヘッダを流用することで出力しています．
• ProgressMeter.jlのprogress_pmap()が便利らしいですが，なぜかうまく動きませんでした．
• Juliaでは画像輝度値は0~1の実数値で扱います．

参考

1. 「高速な計算にむけて」 http://www.cas.cmc.osaka-u.ac.jp/~paoon/misc/julia/post/fast-computing/
2. 「Juliaで超単純にマルチプロセス」 https://qiita.com/kojix2/items/52c0b9cf9c130eae387d