なぜcuDNNのConvolutionは高速なのか
cuDNNはNVIDIAが公開しているDeep Learning用のライブラリである。このライブラリを使うとCaffeやChainerなどのDeep Learning用のソフトウェアの速度が向上する。 この速度向上に寄与している最も大きな部分がConvolutionの高速化である。
個人的には、CPUでどこまで高速にConvolutionが実現できるのかに興味がある。本記事は、その準備段階として、どういう高速化戦略がありえるのかを調べたものである。
Convolutionとは
Convolutionは、日本語では畳み込みと呼ばれる操作である。畳み込み操作自体は何次元のデータ構造に対しても定義できるが、以下では、画像処理でよく使われる、二次元のConvolutionのみを考える。
何も考えずに普通にConvolutionを実装すると、以下の擬似コードのようになるだろう。ただし、簡単のため、境界条件、padding、strideなどは考えない。
for w in 1..W
for h in 1..H
for x in 1..K
for y in 1..K
for m in 1..M
for d in 1..D
output(w, h, m) += input(w+x, h+y, d) * filter(m, x, y, d)
この擬似コードについては、Convolution in Caffe: a memoから借用した。
工夫のしどころはどこにあるのか
擬似コードを見れば分かる通り、Convolutionというのは大半が積和計算であり、この積和計算の回数自体を減らすことはできない難しい。しかし、様々な高速化の余地がある。
- メモリアクセスパターンを改善する。単純に実装すると、キャッシュミスが頻発するので全然性能が出ない。ここが一番大きな問題である。
- AVX命令やSSE命令などを使って、同一クロック内で処理できる積和計算の数を増やす。GPUを使うというのもここの部分の工夫に分類できるか。
- 積和計算以外の命令実行に使われる時間を削減する。例えば、カーネルサイズが小さければ、ループアンローリングが有効である。
もっとも重要なのがメモリアクセスの話で、これをどうにかしないと、その他の工夫は焼け石に水程度にしかならない。 2番目、3番目はたんに実装を頑張りましょうという話で、非常に大変ではあるのだが、ここで語れることはなにもない。そのため、以下では、メモリアクセスパターンについての話題についてのみ取り上げる。
普通に頑張ってランダムアクセスを減らす
最も直感的な高速化手法は、普通に頑張ってランダムアクセスを減らすことである。これに関しては、最内ループからはじめる深層学習(waifu2xの高速化)をみるのがわかりやすいだろう。
一般論として、メモリアクセスについては、以下のような点が重要である。
- ランダムアクセスを減らす。
- ランダムでないにせよ、離れたアドレスへのアクセスを避ける。
- 同じアドレスに対する（メインメモリからの）ロード回数をできるだけ減らす。
言い換えると、メインメモリにはできるだけシーケンシャルにアクセスしろ、可能ならメインメモリからのロード回数も減らせ、ということである。
最内ループからはじめる深層学習(waifu2xの高速化)では、次のような工夫をしている。
- 入力画像のデータ構造が(c, h, w)になっているのを、(h, w, c)の順に並べ替えている。チャンネルを内側に持ってくることで、メモリアクセスの局所性を高めることができる。
- ループの順序をh, w, c, kにする。（ただし、kは出力側のフィルタを指すものとする）
フィルターサイズが3x3の場合は、1座標に対する処理をする間、入力画像側でのキャッシュミスは3回しか起きないし、フィルター側のアクセスはシーケンシャルになるので、たしかにこれで高速になりそうだし、実際、高速になっている。
行列積に変形する
行列積に変形するのは、Loweringとかim2colとか呼ばれている方法である。
Convolutionのコードはちょっとややこしく見えるが、以下のように簡単化することができる。
- フィルタを2次元から1次元に直す。フィルタが複数あることを考えると、フィルタは行列になる。
- フィルタと内積計算をするだけで出力値が計算できるように、入力側の画像も（冗長な）行列に変形する。
画像とフィルタが両方とも行列になるので、後はBLASなりなんなりで行列積を計算すればよい。
この方式の利点は、ミニバッチサイズさえ十分に大きければ、それなりによいパフォーマンスが得られる点である。普通にがんばる方式は、ケースバイケースで、どういうアクセスパターンがよいのか、場合分けして実装しないとよいパフォーマンスが得られない。（その点、waifu2xは3x3のフィルターサイズしかないので、普通にがんばる方式がやりやすい訳である。）
弱点は、入力画像を変換した行列がかなり大きくなる点である。私の理解が正しければ、この方式ではフィルタのサイズと同じだけの冗長なピクセルが必要になるので、フィルタサイズが3x3であれば元の画像サイズの9倍、5x5であれば25倍、7x7であれば49倍のメモリが必要になる。
cuDNNでのconvolutionの実装
cuDNN: Efficient Primitives for Deep Learningによれば、cuDNNのConvolutionの基本は、上記のloweringである。しかし、loweringをそのまま実装すると、メモリ消費量の問題がある。そこで、cuDNNはタイリングとloweringを組み合わせてconvolutionの実装としている。
入力画像をいくつかのタイルに分割してからConvolutionを行うことには、以下の様な利点がある。
- メモリ消費量を削減できる。
- タイルを作っている間に積和計算を実行できる。
- タイルにした方がキャッシュにも乗りやすくなるかも？（これは自信がない）
どういうタイルサイズに分割すればいいかは自明な問題ではないが、その辺りがどうなっているかは上記の論文には書いていない。
maxDNNでのconvolutionの実装
maxDNN: An Efficient Convolution Kernel for Deep Learning with Maxwell GPUs は、cuDNNよりも高速なconvolutionの実装である。ebayのgithubアカウントで公開されているが、なぜかあまり他のところで見かけたことがない。
maxDNNのセールスポイントは、フィルタサイズが3x3〜11x11程度の場合に、90%以上の演算効率を発揮する点である。フィルタサイズが3x3であれば、95%程度の演算効率になる。convolutionをそのまま実装するだけでは、これ以上速くするのはかなり難しいだろう。
maxDNNのconvolutionは、以下のような特徴がある。
- cuda-convnet2スタイルのconvolutionを使う。すなわち、入力画像の形式を(h,w,c,n)の順番にする。これはメモリアクセスの局所性を上げるためである。
- maxasというCuda用アセンブラのサンプルとして付属している、高速なSGEMM（行列積）を使って行列積を計算する。
SGEMMを使っているということは、loweringを使っているか、もしくはcuDNNのようなtiling + loweringのどちらかだろう。
Origami: A 803 GOp/s/W Convolutional Network Accelerator によれば、NervanaSystemsのneon（とにかくconvolutionが速いことで有名）もmaxDNNベースらしい。maxasを作っているのはNervanaSystemsのScott Grayだし、maxDNNの著者Andrew LavinはScott Grayと共著で論文を書いてるし、この辺りの人々は関係が深いようだ。
その他の高速化技法
Convolutionの定義に従ったまま高速化するのであれば、おそらく、maxDNNの効率を超えることは難しい。しかし、近似してもよいのであれば、さらなる高速化が可能である。面白いことに、近似したからといって精度が必ず下がるわけではない。以下に、いくつか見つけたものを挙げてみる。
- Separableなフィルタを使うことで積和計算の回数を減らす
- Winograd Minimal Filtering Algorithmとやらを使う
- 行列積とみなした時に、フィルタ側の行列を低ランク近似する
低ランク近似はわかりやすくて気に入っている。Faster R-CNNの人たちも、fully connected layerに対して似たようなことをしている。
今後についての予想
直感的には、waifu2xの高速化で出てきたデータの並べ替えは十分に効率が良さそうに思えていたのだが、実際、maxDNNは似たような手法で高い性能を発揮しているわけで、この方法が（速度的には）正しいことは既に証明されていると言えるだろう。
ではなぜ他のライブラリでは似たような手法を取らないのだろうか。バッチ次元を一番上から動かしたくないのでは、というのが私の予想である。コードを書く上では、バッチ次元は一番最初に来た方がわかりやすいので、多くのレイヤーが、バッチ次元を最初に持ってくる形で実装されてしまっている、そうなると、他のレイヤーとつなぐ上で、バッチ次元を動かすのは得策ではない。
しかし、convolutionは明らかに計算のボトルネックになっているし、Deep Learningにおいて、学習時間は大きな問題となっている。近い将来、convolutionの実装はmaxDNN方式に向かうだろう。
私の知る限り、Chainer以外のDeep Learningライブラリは、ネットワークを表現するグラフを定義する、という方向に進化している。既存のconvolutionをpreordering, convolution_maxdnn, postorderingみたいな感じに分割してやって、連続したconvolutionの部分はpostordering, preorderingでキャンセルがかかる、みたいな処理が入れば、バッチ次元の移動回数は最小限に絞ることができるだろう。そういう風に、グラフ全体を見て処理を最適化する、みたいな方向に進むと予想している。
感想
当初はcuDNNについて調べていたのだが、最終的には、maxDNNがすごいという結論になってしまったが、Convolutionを高速にする上で、なにかとてつもない銀の弾丸があるわけではないことがわかったのは良かった。
convolutional neural networkはかなり歴史が深いモデルなのに、その基礎となるconvolution操作は、近年まであまり最適化の対象としてみなされていなかったような雰囲気を感じる。それを意外に感じるとともに、一度重要だとみなされれば、あっという間にGPUの理論限界近くまで性能が引き出されるところに、今の研究スピードの速さを改めて思い知らされた。
