CuPyを使った二次元FFTの高速化

先日のGTC2018でnumpyのFFTがCupyで動くということを知りました

以前、numpyで二次元FFTをやっていて遅かったので、どのくらい改善するのかトライしてみました

結論から言うと、データが大きいかCPUがしょぼい場合はGPUを使った方が早いです

機材

　1)8086K + GTX1070
　2)6500 + GTX960
　3)Jetson TX2

環境

機材1),2)
　OS:Windows10
　GPUのドライバー：399.07
　CUDA:9.0
　CuDNN:v7.1
　python:3.6.6
　CuPy:4.1.0
　numpy:1.15.1

機材3)
　Jetpack：3.3
　python:3.5.2
　CuPy:4.4.1
　numpy:1.15.0

numpy

import CV2
import numpy as np

#画像の読み込み
img = cv2.imread(image, 0)
#リサイズ
img = cv2.resize(img, (200,200))
#float32に変換
img = img_fft.astype(np.float32)

#2次元FFT処理
f = np.fft.fft2(img)
fshift = np.fft.fftshift(f)
ms = 20*np.log(np.abs(fshift))

#uint8に戻す
ms = ms.astype(np.uint8)

CuPy

GPUへの行きと帰りを書くだけでOKです

import CV2
import numpy as np
import Cupy as cp

#画像の読み込み
img = cv2.imread(image, 0)
#リサイズ
img = cv2.resize(img, (200,200))
#float32に変換
img = img_fft.astype(np.float32)

#GPUに送る
f = cp.asarray(img)

#####ここからGPU#####
#2次元FFT処理
f = cp.fft.fft2(f)
fshift = cp.fft.fftshift(f)
ms_gpu = 20*cp.log(cp.abs(fshift))
#uint8に戻す
ms_gpu = ms.astype(cp.uint8)
#####ここまでGPU#####

#CPUに戻す
ms_cpu = cp.asnumpy(ms_gpu)

性能比較（1回目）

200pix,200pixの画像960枚を変換

CPU	GPU	numpy(sec)	CuPy(sec)	%
8086K	GTX1070	0.78	1.72	220%
6500	GTX960	0.91	1.93	212%
Jetson	TX2	8.11	3.32	-60%

キャッシュに乗るくらいだとCPUの方が早い模様
Jetsonはこの時点でかなりの高速化

性能比較（2回目）

800pix,800pixの画像960枚を変換

CPU	GPU	numpy(sec)	CuPy(sec)	%
8086K	GTX1070	12.44	4.80	-61%
6500	GTX960	14.19	8.09	-43%
Jetson	TX2	126.05	20.63	-84%

このくらいになってくるとPCでもGPU使用が圧勝