Motivation
FCNN: Fourier Convolutional Neural Networks(http://ecmlpkdd2017.ijs.si/papers/paperID11.pdf )を読んでconvolutionをフーリエドメインで行うことに興味が湧いたのでnumpyで実際に挙動を確認してみました.
空間ドメインでconvolutionをnaiveに実装すると4重ループが出てきてこの計算は結構しんどいのですが, この演算をフーリエドメインで行うとただの要素積になります(スカラーなら通常の積, 行列はアダマール積). フーリエ変換&逆変換は線形変換なので結局ループ計算は出てくるのですが、対称性のある行列積なので、高速な計算方法が知られており、空間ドメインでのconvolutionの計算のループより高速に動作します. そこで、入力と畳み込みカーネルを共にフーリエ変換してフーリエドメインで積を計算し、その後入力をフーリエ逆変換で空間ドメインに戻すとnaiveなconvolutionの実装より高速に動きます(入力が大きいほど顕著に効果が出る). 論文中ではどのくらい高速化するかなども書いてありますが, 今回は実装そのものが目的なので書きません.
Preliminary
フーリエ変換とconvolutionの関係
二次元の離散フーリエ変換(DFT)に関してはwikipediaがかなり詳しいです.
deepの文脈におけるconvolutionは数学的には畳み込み(convolution)ではなく, 相関関数(cross-correlation)と呼ばれています. フーリエドメインでの相関関数の計算は,二つの入力(画像と畳み込みカーネル)のうち, 一方の複素共役をとったものとの間で要素積をとります.
$$z(\tau)=\int^{\infty}_{-\infty}h(t)x(t+\tau)dt=\mathcal{F}^{-1}[\mathcal{F}(h)\circ\mathcal{F}^{*}(x)](\tau)$$
padding
deepの実装では入力の空間方向へのサイズを一定の大きさに保つためにしばしばpaddingという操作を行います。
今回の実装ではreflective paddingとzero paddingの2種類を使います. reflective paddingは入力に周期性を仮定して, convolutionが入力のエッジに到達したら逆サイドのピクセルを畳み込みます. zero paddingは入力の外側のピクセル値を全て0だと思って畳み込みます. chainerなど深層学習ライブラリでは基本的にzero paddingが使われている気がします. そもそも入力に周期性はないですし. reflective paddingはその意味で"フーリエっぽい"です. フーリエ変換を使う解析ではこちらの方が使われたり使われなかったりします.
その他
離散フーリエ変換は本来連続なフーリエ変換を離散化しているので, 変換(逆変換)の過程で数値誤差が発生することがあります. 実装ではこの辺りは適当に丸めて誤魔化しています. また, reflective paddingに伴う周期ずれは適当に位相を合わせています.
Experiment
まず2種類のpaddingを実装します.
直感的な理解を助けるためにそれぞれの挙動を貼っておきます.
import cmath
import numpy as np
def zero_padding(M, pad=1):
canvas = np.zeros((M.shape[0] + 2*pad, M.shape[1] + 2*pad))
canvas[pad:-pad, pad:-pad] = M
return canvas
def reflective_padding(M, pad=1):
num_tile = 2 * ((pad - 1) // M.shape[0] + 1) + 1
M_tile = np.tile(M, (num_tile, num_tile))
crop_index = int((M_tile.shape[0] - (M.shape[0] + 2*pad))/2)
return M_tile[crop_index:-crop_index, crop_index:-crop_index]
# reflective padding
f = np.arange(25).reshape(5, 5)
reflective_padding(f)
# 戻り値
array([[24, 20, 21, 22, 23, 24, 20],
[ 4, 0, 1, 2, 3, 4, 0],
[ 9, 5, 6, 7, 8, 9, 5],
[14, 10, 11, 12, 13, 14, 10],
[19, 15, 16, 17, 18, 19, 15],
[24, 20, 21, 22, 23, 24, 20],
[ 4, 0, 1, 2, 3, 4, 0]])
# zero padding
zero_padding(f)
# 戻り値
array([[ 0., 0., 0., 0., 0., 0., 0.],
[ 0., 0., 1., 2., 3., 4., 0.],
[ 0., 5., 6., 7., 8., 9., 0.],
[ 0., 10., 11., 12., 13., 14., 0.],
[ 0., 15., 16., 17., 18., 19., 0.],
[ 0., 20., 21., 22., 23., 24., 0.],
[ 0., 0., 0., 0., 0., 0., 0.]])
convolution, DFT版convolutionの実装
# フーリエドメインを経由した畳み込みの結果が空間ドメインでの畳み込みの結果に一致することを確認するため,
# naiveな畳み込みも実装しておきます
def convolution_space(M, kernel):
out_h = M.shape[0] - kernel.shape[0] + 1
out_w = M.shape[1] - kernel.shape[1] + 1
out = np.zeros((out_h, out_w))
for i in range(out_h):
for j in range(out_w):
out[i, j] = np.sum(M[i:kernel.shape[0]+i, j:kernel.shape[1]+j] * kernel)
return out
# 入力(正方行列)の一辺の大きさを引数に取って, 離散フーリエ変換&逆変換, フーリエドメインでの畳み込みを
# 実行できるクラスを実装します
class DFT:
def __init__(self, input_shape):
self.M = self.dft_matrix(input_shape)
self.ctM = np.conjugate(self.M.T)
def dft_matrix(self, n):
A = np.arange(n)
T = A.reshape(1, -1)
X = A.reshape(-1, 1)
M = cmath.e**(-1j * 2 * cmath.pi * T * X / n)
return M
def dft(self, f):
return self.M @ f.T @ self.M
def idft(self, F):
return (self.ctM @ F @ self.ctM).T.real / np.prod(F.shape)
def conv(self, f, kernel):
FT_f = self.dft(f)
FT_kernel = np.conjugate(self.dft(kernel))
out = FT_f * FT_kernel
return self.idft(out)
フーリエ変換を介したconvolution
まずはreflective paddingを使った場合の畳み込みをフーリエ変換を使って計算し, 空間ドメインでの畳み込みと結果が一致することを確認します.
# reflective padding
kernel = np.arange(9).reshape(3, 3)
# 入力fとサイズが揃うようにkernelをzero paddingする
kernel_expand = np.zeros_like(f)
h_k, w_k = kernel.shape
kernel_expand[:h_k, :w_k] = kernel
kernel_expand
# 空間ドメインにおける拡大した畳み込みカーネル
array([[0., 1., 2., 0., 0.],
[3., 4., 5., 0., 0.],
[6., 7., 8., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]])
# DFTによるconvolution
transform = DFT(5)
out = transform.conv(f, kernel_expand).astype('int32') # 確認のため数値誤差を整数化して誤魔化す
np.roll(out, 1, [0, 1]) # ズレた周期を補正
# 戻り値
array([[216, 207, 243, 279, 240],
[321, 312, 348, 384, 345],
[501, 492, 528, 564, 525],
[681, 672, 708, 744, 705],
[336, 326, 363, 399, 360]], dtype=int32)
# 結果の確認
convolution_space(reflective_padding(f), kernel)
# DFTの結果と一致
array([[216., 207., 243., 279., 240.],
[321., 312., 348., 384., 345.],
[501., 492., 528., 564., 525.],
[681., 672., 708., 744., 705.],
[336., 327., 363., 399., 360.]])
周期がずれたことをのぞいて結果が一致しました.
# 入力fとサイズが揃うようにkernelをzero paddingする
# zero paddingの場合補完した部分を含めた入力が1周期に相当するので, 入力サイズは(7, 7)
kernel_expand = np.zeros((7, 7))
h_k, w_k = kernel.shape
kernel_expand[:h_k, :w_k] = kernel
kernel_expand
# 結果
array([[0., 1., 2., 0., 0., 0., 0.],
[3., 4., 5., 0., 0., 0., 0.],
[6., 7., 8., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0.]])
# DFTによるconvolution
transform = DFT(7)
out = transform.conv(zero_padding(f), kernel_expand).astype('int32') # 確認のため数値誤差を整数化して誤魔化す
out[:5, :5] # zero paddingによって増えた次元を削る
# 戻り値
array([[ 88, 141, 174, 208, 135],
[209, 311, 347, 383, 239],
[344, 491, 527, 563, 344],
[480, 671, 707, 743, 449],
[232, 303, 319, 333, 184]], dtype=int32)
# 結果の確認
convolution_space(zero_padding(f), kernel)
# DFTの結果と(ほぼ)一致
array([[ 88., 142., 175., 208., 136.],
[210., 312., 348., 384., 240.],
[345., 492., 528., 564., 345.],
[480., 672., 708., 744., 450.],
[232., 304., 319., 334., 184.]])
だいたい一致しました. 微妙にズレがあるのはDFTによる変換→逆変換の数値誤差かと思われます.
振り返り
今回実装で困った点がありました.
それは**入力とkernelのサイズが違う時はどうすればいいか(アダマール積をとるので同じサイズである必要があるが一般にkernelは入力よりずっと小さい)**という問題です.
これは入力と同じサイズになるようkernelをzero paddingすることで解決しました. 畳み込みカーネル1つが入力画像1枚のキャンバスに対応するのでそれはそうかなという感じがしますがめちゃくちゃハマりました(ちなみにzero paddingする場所はどこでも結果はほぼ変わりません. 周期がずれるだけです).
今回はDFTをスクラッチで実装しましたが, おそらくnp.fft.fft2やnp.fft.ifft2を使った方が早かったり, 数値誤差が少なかったりしそうです.