概要

前回の記事で述べましたが、FCNNの畳み込みは非常にメモリ効率が悪いことが知られています。
そこでOverlap-Add: OVAまたはOLA: 重畳加算法の導入が提案されました。これによりフィルタサイズと入力サイズとを揃える必要がなくなり、メモリ効率が大幅に向上しました。
この手法はもともと信号処理での畳み込み演算に利用されている手法ですが、同じ畳み込み(仮)ですのでCNNにも適用することができます。
本記事ではそのOVA法と、ついでにOVA法と類似のOverlap-Save: OVSまたはOLS: 重畳保留法について触れています。

何かの役に立てばぜひLGTM、ストック、コメントしていただけると励みになります。

Overlap-Add法

OVerlap-Add: OVA法はもともと信号処理の分野で非常に長い入力信号に対して随時畳み込みをかけることで結果的に高速に全体を畳み込みするための手法です。まずは1次元でのアルゴリズムを見てみましょう。
入力の長さ$I$、フィルタサイズ$F$とすると

入力をブロック長$F$のブロックに重複なく分割する
フィルタと分割ブロックとを線形畳み込みする
各ブロックを$F-1$だけ重ねて、重なっている部分を足し算しつつブロックを繋げる

となります。シンプルですね〜
実際に図を交えながら詳細に見ていきましょう。

入力を$(x_1, x_2, \ldots, x_9)$、フィルタを$(w_1, w_2, w_3)$とします。
まずは入力をフィルタと同サイズのブロックに分割します。

この分割それぞれとフィルタとを線形畳み込みで畳み込みます。

こうして計算された出力を、$F-1$だけ重ね合わせて配置し、重なっている部分は足し算することで最終的な出力を得ることができます。

しっかり線形畳み込みの結果になっていますね！

Overlap-Save法

続いてOverlap-Save: OVS法について紹介します。こちらはちょっとわかりづらい＆あまり解説しているサイトが見当たらなかったので間違っている部分もあるかもしれません...

先ほどと同じように定数を設定します。また、ブロック長を$N$とします。この$N$は任意数です。ただし$N$は入力長$I$の倍数になっている方がパディングする必要もなくていいかもしれません。
各ブロック$F-1$だけ重ねてブロック長$N$となるように分割します。

続いてこの各ブロックとフィルタとを循環畳み込みで畳み込みます。

このように計算された出力の、$F-1$だけ破棄したものを結合することで最終的な出力を得ることができます。

CNNに適用する

ここまででOVA法、OVS法の1次元での動作を見てきました。ここからはCNNに適用することを考えます。
といっても、ただ単に2次元に拡張してあげればOKです。
また、色々と調べている中でFCNNにおいてOVS法を使用している例は見つからなかったため、OVA法のみ載せます。

まあシンプルな線形畳み込みです。CNNでの畳み込みとは微妙に異なりますが、周囲$(F_h-1, F_w-1)$だけカットすればCNNでの畳み込みと同等になりますね。もちろん添字を反転させる必要もあります。詳しくは前回の記事参照。

FCNNに適用する

さて、Fourier CNNに適用しましょう。どうやって適用するのか、色々な論文などを流し読みしてみたりしましたが、どれもこれも微妙...流し読みしてるせいかもしれませんが笑
ということで自分で色々考えて試しに手計算してみたりしたのですが、結局以下のようにするしかなさそうという結論になりました。

図中のFFT,IFFTは高速フーリエ変換と高速逆フーリエ変換を意味しています。
また、流石にきちんと図中に記述するのは無理だったのでフーリエ変換結果などは省略して大文字で表してあります。
メモ代わりに少しだけ計算式を載せておきます。

X^{(0)}_{m,n} = \sum_{k=0}^{F_h-1}{\sum_{l=0}^{F_w-1}{x_{k,l}e^{i\frac{2\pi}{F_h}km}e^{i\frac{2\pi}{F_w}ln}}} \\
W_{m,n} = \sum_{k=0}^{F_h-1}{\sum_{l=0}^{F_w-1}{w_{k,l}e^{i\frac{2\pi}{F_h}km}e^{i\frac{2\pi}{F_w}ln}}} \\
y^{(0)}_{m,n} = \frac{1}{F_h}\sum_{k=0}^{F_h-1}{\frac{1}{F_w}\sum_{l=0}^{F_w-1}{X^{(0)}_{k,l}W_{k,l}e^{-i\frac{2\pi}{F_h}km}e^{-i\frac{2\pi}{F_w}ln}}}

このように分割した部分ごとにフーリエ変換→要素積→逆フーリエ変換とした後でOVA法に則って計算していけば、省メモリを実現しつつ求めたい畳み込み演算結果を得ることができるわけです。

計算量の比較

ただの線形畳み込み
OVA法を利用した線形畳み込み
高速フーリエ変換を利用した線形畳み込み
OVA法と高速フーリエ変換を利用した線形畳み込み

について、それぞれ時間計算量・空間計算量を調べてみます。
ただし、簡略化のために入力、フィルタ、出力の形状は全て正方形であるとします。

なお様々な参考資料を独自解釈して考えているため、間違いがある可能性があります。
また、オーダ記法$\mathcal{O}$はできる限り採用していません。

1. ただの線形畳み込み

時間計算量：$2O^2F^2$
空間計算量：$I^2+F^2+O^2$

ただの線形畳み込みは以下の数式で表すことができます。

y_{m,n} = \sum_{k=0}^{F-1}{\sum_{l=0}^{F-1}{x_{m-k,n-l}w_{k,l}}}

この部分の時間計算量は$F^2$回の足し算と$F^2$回の掛け算であり、これが出力サイズの回数分だけ実行されるため、$2O^2F^2$が最終的な時間計算量となります。
ストライド1、パディング0の場合なら$O=I-F+1$などと計算できますね。
空間計算量については、もともとの入力$I^2$と重み$F^2$、出力の$O^2$を足し合わせたものになります。

2. OVA法を利用した線形畳み込み

時間計算量：$\frac{I^2}{F^2}-1 + 2(F+1)^2F^2 + 2 \left( \frac{I}{F}-1 \right)O(F-1) + \frac{1}{2}\frac{I}{F} \times \frac{1}{2}\frac{I}{F}(F-1)^2$
空間計算量：$I^2+F^2+O^2$

OVA法についての正確な計算量の導出は不明なのですが、とりあえず独自解釈を交えて考えてみます。
間違っていればご指摘ください。

まず、アルゴリズムを再掲します。

入力をブロック長$F$のブロックに重複なく分割する
重みと分割ブロックとを線形畳み込みする
各ブロックを$F-1$だけ重ねて、重なっている部分を足し算しつつブロックを繋げる

一つずつ考えてみましょう。ただし、簡略化のためにパディング不要などの色々といい感じの条件下でのことにします。

入力を分割する回数は$$\frac{I^2}{F^2}-1回$$生成されるブロックのサイズは$F \times F$
それぞれを線形畳み込みするのに必要な計算回数は$$2(F+1)^2F^2回$$生成される出力ブロックは$o \times o$
ブロックの数は$\frac{I}{F} \times \frac{I}{F}$個であり、先の図で言う色の濃い帯は$2 \times \left( \frac{I}{F}-1 \right)$本、帯の幅が$F-1$となるため、足し算の回数は$$2 \left( \frac{I}{F}-1 \right)O(F-1)回$$ただし、$2\times2$個のブロックセットにつき$(F-1)(F-1)$回ずつ、数えられていない足し算が存在するので、それら$\frac{1}{2}\frac{I}{F} \times \frac{1}{2}\frac{I}{F}(F-1)^2$回も考慮した$$2 \left( \frac{I}{F}-1 \right)O(F-1) + \frac{1}{2}\frac{I}{F} \times \frac{1}{2}\frac{I}{F}(F-1)^2回$$

よって、時間計算量はこれら全てを足し合わせた

\frac{I^2}{F^2}-1 + 2(F+1)^2F^2 + 2 \left( \frac{I}{F}-1 \right)O(F-1) + \frac{1}{2}\frac{I}{F} \times \frac{1}{2}\frac{I}{F}(F-1)^2

となります。
空間計算量については、あらかじめ出力サイズ$O \times O$のメモリを確保しておき、ブロックごとの線形畳み込みによる出力ブロックをこのメモリの該当箇所に足し合わせていく、と考えれば$O^2$となります。
もちろん入力、重みの空間計算量$I^2$と$F^2$も加算します。

3. 高速フーリエ変換を利用した線形畳み込み

時間計算量：$2I^2\log_2I + 2F^2\log_2F + 2O^2\log_2O + 6O^2$
空間計算量：$6O^2$

高速フーリエ変換を利用する場合、入力やフィルタを出力サイズに拡張しつつフーリエ変換を行う必要があります。このフーリエ変換に必要な時間計算量は、高速フーリエ変換の場合は入力、重みそれぞれについて$2I^2\log_2I$と$2F^2\log_2F$となります。
そして入力と重みとの要素積を取るのですが、この時の要素積は複素数積であることに注意すると$$(a+bi)(c+di)=ac-bd+i(ad+bc)$$より積4回和1回差1回となり、通常の6倍の演算量であると見なすことができるので、$6O^2$となる。
最後に出力を逆フーリエ変換する必要があるため、その計算量$2O^2\log_2O$が加えられます。
空間計算量については、入力及び重み、出力の3つがそれぞれ$O \times O$必要で、さらに複素数として保持する必要があるため実数と比較して倍のメモリが必要となり、$6O^2$となります。

ちなみに高速フーリエ変換の時間計算量は厳密な計算量の出し方がわからなかったので、厳密なものではなくオーダです。

4. OVA法と高速フーリエ変換を利用した線形畳み込み

時間計算量：$\frac{I^2}{F^2}-1 + 2 \left( \frac{I^2}{F^2}+1 \right)F^2\log_2F + 6\frac{I^2}{F^2}o^2 + 2\frac{I^2}{F^2}o^2\log_2o$
空間計算量：$I^2 + 2 \left( 1 + \frac{I^2}{F^2} \right)o^2 + O^2$

OVA法を利用する場合のアルゴリズムは以下のようになります。

入力をブロック長$F$のブロックに重複なく分割
各入力ブロックと重みを高速フーリエ変換
重みと入力ブロックとを要素積
生成された出力ブロックを逆フーリエ変換
各ブロックを$F-1$だけ重ねて、重なっている部分を足し算しつつブロックを繋げる

1と5は先の議論から$\frac{I^2}{F^2}-1$回と$2 \left( \frac{I}{F}-1 \right)O(F-1) + \frac{1}{2}\frac{I}{F} \times \frac{1}{2}\frac{I}{F}(F-1)^2$回ですね。それ以外の部分を見ていきます。

各入力ブロックと重みはサイズ$F \times F$なので、高速フーリエ変換の計算量は$2F^2\log_2F$であり、全部で$\frac{I}{F} \times \frac{I}{F} + 1$個あるため$$2 \left( \frac{I^2}{F^2}+1 \right)F^2\log_2F$$となります。また、生成されるブロックサイズを$o \times o$としておきます
重みと入力ブロックの要素積に必要な計算回数は、複素数積であることを考慮すると$6o^2$であり、これが$\frac{I}{F} \times \frac{I}{F}$回行われるので$$6\frac{I^2}{F^2}o^2回$$
出力ブロックのサイズは$o \times o$で、それが$\frac{I}{F} \times \frac{I}{F}$個あるため$$2\frac{I^2}{F^2}o^2\log_2o$$

以上から、時間計算量は総じて

\frac{I^2}{F^2}-1 + 2 \left( \frac{I^2}{F^2}+1 \right)F^2\log_2F + 6\frac{I^2}{F^2}o^2 + 2\frac{I^2}{F^2}o^2\log_2o

となります。
空間計算量については、入力の$I^2$と出力$O^2$はもちろん、重みや入力ブロックごとに$2o^2$必要となります。

おわりに

GIFアニメーションの手抜きっぷりよ...気が向いたらもう少しマシなの作ります。

CNN高速化シリーズ

CNN高速化シリーズ ~FCNN: Fourier Convolutional Neural Networkの導入~

CNN高速化シリーズ ~Overlap-Add法とOverlap-Save法~

概要

目次