Chainerで書いたNNとダミー入力があるとき、そのNNのforward passの理論的な計算量・メモリ転送量を計算するchainer_computational_costを作りました。
ChainerのFunction Hookをベースにしているため、NNの定義コードに手を入れる必要は一切ありません。
背景
TL;DR 闇雲に軽いNNを探すより定量的で気軽に使える指標がある方がいいですよねという話です。
ニューラルネット技術の実用化が進む中でいかに精度を保ちながら軽いアーキテクチャを探すかはとても重要です。手動にせよ自動にせよ、基本的には設計→学習→計測という試行錯誤のサイクルを回すのが一般的と思います。実際には、誰しもが当然「こういうアーキテクチャのNNなら理論的には倍ぐらい軽いはずだ」という仮説、すなわち理論計算量を頭の中で相対的に概算してNNを設計していることでしょう。この理論計算量が絶対的な数字として利用可能であれば、サイクルを回すのがとても効率的になります。
もちろん理論計算量と実際にプロセッサ上で走る計算の回数は一般には異なる上、単位理論計算量に対する実際の速度は演算の種類によっても異なります。したがって最終的な推論パフォーマンスは常に実機上でのベンチマークに基づいて議論するべきです。一方で常に最速のサイクルで設計→学習→デプロイ→計測をすることが物理的には困難なこともままあるため、枝刈り的に理論計算量を活用することは依然重要です。少なくとも、パラメータ数によってNNの重さを議論するよりはずっと意味があります。
また理論計算量がわかると、実機のカタログスペック上の性能をつかって理論上の推論速度を概算することができます。これを実機上でのベンチマークの結果と突き合わせることで性能の実効効率の概算を求められます。これは実装の最適化を議論するときの手がかりとして非常に有用です。闇雲な目標値に向かって先の見えない最適化をしていくよりも意味のある議論ができます。
できること
Layerごとに、次のような項目を集めることができます。
- 理論計算量(Number of floating point operations (FLOPs(*1) ))
- 理論メモリ読み込み量(bytes)
- 理論メモリ書き込み量(bytes)
これらを、オペレータの種類ごとに集計すること、またNN全体として集計することができます。
これをCSVやmarkdown、テキストテーブルの形で出力できます。
基本的な仕組み
ChainerにはFunctionHookという仕組みがあり、Function（NNにおける層と同義と思ってOK）が実行される度に走る処理というのを定義することができます。そのときのFunctionのオブジェクトの保持するハイパーパラメータ（Convolutionのkなど）、入力のshapeを使うことで、層ごとの理論計算量を計算することができます。
オペレータごとに計算式を用意して、hook関数が受け取ったFunctionオブジェクトの型によって適切な計算式を選んで理論計算量を求めて集計しているというのがchainer_computational_costの本質的な仕組みです。
超単純です。
理論計算量の定義
基本的には、「最も単純でナイーブな実装をしたとき、入力データに対して適用されるfloating point operationの回数」を理論FLOPsとしています。すなわち、添字計算などのような、入力データに対しての計算ではないものは考えていません。メモリ転送量についても、その層が本質的に必要とするデータの読み込み・書き込みが1回だけおこるものとしています。
例えばConvolutionの最適な実装ではim2colのような形でメモリ内で冗長な表現にした上で演算をしたり、メモリアクセス効率向上の為にCHWをHWCに変換したりするなどの処理が行われたりしていますが、これらは一切無視されています。またWinogradなどにより重たい乗算を減らし加算を増やすことで速度を稼ぐことも一般的ですが、ここでは全ての四則演算を同じ1FLOPとしてカウントしています。
加えて、実際のNN推論最適化においてはLayer fusionという最適化が行われます。典型的にはConv-BN-ReLUのようなスタックを1つのCUDAカーネルで行うことで、特にメモリ転送量を大きく減らす効果があります（カーネル呼び出しオーバヘッドを削減する意図もあります）。理論計算量ではそういった要素も無視され、全てのlayerがデータ読み込み→演算→データ書き込みとする想定で計算されています。
今のところ、オペレータごとにその「最もナイーブな実装」として何を想定するかは結構ふわっとしていて、わりと雑なところもあります。そもそものFLOPsの計算式も誤っているところもおそらくあるので、しばらくはそのへんの修正が入ることと思います。
導入と活用例
pipで入れられます。
% pip install chainer_computational_cost
ResNet50の理論計算量を出してみましょう。本質的には2行足すだけです。
import numpy as np
import chainer
import chainer.links as L
import chainer_computational_cost as c3
net = L.ResNet50Layers()
s = (1, 3, 224, 224)
x = np.random.random(s).astype(np.float32)
with chainer.no_backprop_mode(), chainer.using_config('train', False):
with c3.ComputationalCostHook(fma_1flop=True) as cch:
y = net(x)
cch.show_report(unit='M', mode='md')
show_reportによって、層ごとの理論計算量をレポートできます。
| Layer name | MFLOPs | MemRead MiB | MemWrite MiB | MemR+W MiB |
|---|---|---|---|---|
| Convolution2DFunction-1 | 118.014 | 0.61 | 3.062 | 3.673 |
| FixedBatchNormalization-1 | 0.803 | 3.063 | 3.062 | 6.125 |
| ReLU-1 | 0.803 | 3.062 | 3.062 | 6.125 |
| MaxPooling2D-1 | 1.606 | 3.062 | 0.766 | 3.828 |
| Convolution2DFunction-2 | 12.845 | 0.781 | 0.766 | 1.547 |
| FixedBatchNormalization-2 | 0.201 | 0.766 | 0.766 | 1.532 |
| ReLU-2 | 0.201 | 0.766 | 0.766 | 1.531 |
| Convolution2DFunction-3 | 115.606 | 0.906 | 0.766 | 1.672 |
| FixedBatchNormalization-3 | 0.201 | 0.766 | 0.766 | 1.532 |
| ReLU-3 | 0.201 | 0.766 | 0.766 | 1.531 |
| Convolution2DFunction-4 | 51.38 | 0.828 | 3.062 | 3.891 |
| FixedBatchNormalization-4 | 0.803 | 3.064 | 3.062 | 6.127 |
| Convolution2DFunction-5 | 51.38 | 0.828 | 3.062 | 3.891 |
| FixedBatchNormalization-5 | 0.803 | 3.064 | 3.062 | 6.127 |
| Add-1 | 0.803 | 6.125 | 3.062 | 9.188 |
| ReLU-4 | 0.803 | 3.062 | 3.062 | 6.125 |
| … | ||||
| total | 3884.872 | 256.754 | 137.254 | 394.008 |
ResNet50に画像1枚をかますと、理論的には3884MFLOPs、つまり3.88*10^9回の浮動小数点数演算が実行されることを報告しています。FMA(Fused-Multiply-Add, ax+bの形で表される3入力の演算)については1命令として扱っています（fma_1flop=Trueによって）。実効性能を議論するにはそのほうが都合がいいためです。
またニューラルネットでは層ごとに入力をメモリから読み、また重みをメモリから読み、演算をした上で結果をメモリに書き出すので、演算能力よりメモリ帯域が性能を制約することもよくあります。そのため、理論メモリ転送量の推定も同じようにレポートされています。ResNet50の場合はトータルで394MBのメモリ読み書きが発生することがわかります。
これを元に特定デバイスでの理論性能（の上限）を概算してみましょう。
例えばNVIDIA GTX1080Tiを考えてみると、CUDAコアの演算能力はFP32で10.6TFLOPS、メモリの転送速度が484GB/sです。極めて単純な計算ですが、理論計算量をこのカタログスペックで割ると演算に366マイクロ秒・メモリ転送に814マイクロ秒、合計すると1.2ミリ秒あたりが画像1枚のforward passにかかる最短時間になるだろうと概算できます。
（実際はメモリ転送と演算がオーバーラップできること、Conv->BN->ReLUなどは1つのCUDAカーネルで一気にやるのでメモリ転送量が大幅に減ること、などの要素が絡んできますが、理論計算量に基づく概算の性能上限ということで無視しています）
このように理論計算量から求めた概算性能から、実効効率を計算してみます。
ONNX-Chainerを使ってこのChainerモデルをONNXに吐き出し、ONNX-TensorRTを使ってNVIDIAの推論エンジンTensorRTでResNet50を走らせると、FP32で平均2.41957ms/imgという実効性能が得られました。これは、理論計算量からの概算に対し50%ほどの効率が出ていることを表します。
理論計算量においては、インデックス計算などのような本質的ではないが実際は必要である演算は含まれていないので、それを込みでこの50%という効率は相当のものだと思います。
ミニバッチにも対応していて、ダミー入力を単純にバッチ方向に重ねればOKです。
下記では、ダミー入力によるforward passもバッチサイズ32となると重たいのでGPUで実行しています。
import numpy as np
import chainer
import chainer.links as L
import chainer_computational_cost as c3
net = L.ResNet50Layers()
net.to_gpu()
s = (32, 3, 224, 224)
x = cupy.random.random(s).astype(np.float32)
with chainer.no_backprop_mode(), chainer.using_config('train', False):
with c3.ComputationalCostHook(fma_1flop=True) as cch:
y = net(x)
cch.show_report(unit='G', mode='md')
| Layer name | GFLOPs | MemRead GiB | MemWrite GiB | MemR+W GiB |
|---|---|---|---|---|
| Convolution2DFunction-1 | 3.776 | 0.018 | 0.096 | 0.114 |
| FixedBatchNormalization-1 | 0.026 | 0.096 | 0.096 | 0.191 |
| ReLU-1 | 0.026 | 0.096 | 0.096 | 0.191 |
| MaxPooling2D-1 | 0.051 | 0.096 | 0.024 | 0.12 |
| Convolution2DFunction-2 | 0.411 | 0.024 | 0.024 | 0.048 |
| FixedBatchNormalization-2 | 0.006 | 0.024 | 0.024 | 0.048 |
| ReLU-2 | 0.006 | 0.024 | 0.024 | 0.048 |
| Convolution2DFunction-3 | 3.699 | 0.024 | 0.024 | 0.048 |
| FixedBatchNormalization-3 | 0.006 | 0.024 | 0.024 | 0.048 |
| ReLU-3 | 0.006 | 0.024 | 0.024 | 0.048 |
| Convolution2DFunction-4 | 1.644 | 0.024 | 0.096 | 0.12 |
| FixedBatchNormalization-4 | 0.026 | 0.096 | 0.096 | 0.191 |
| Convolution2DFunction-5 | 1.644 | 0.024 | 0.096 | 0.12 |
| FixedBatchNormalization-5 | 0.026 | 0.096 | 0.096 | 0.191 |
| Add-1 | 0.026 | 0.191 | 0.096 | 0.287 |
| ReLU-4 | 0.026 | 0.096 | 0.096 | 0.191 |
| … | ||||
| total | 124.316 | 5.072 | 4.289 | 9.361 |
同様に124GFLOPsという理論計算量、9.36GBという理論メモリ転送量をGTX1080Tiのカタログスペックにあてはめると、演算に11.7ms、メモリ転送に19.3ms、合わせて31ms/batchというのが理論上のforward passの最短時間として得られます。
実際のバッチサイズ32での推論速度は30.656msでした。
効率が100%を超えてしまったのは、上述のように演算とメモリ転送のオーバーラップや、Conv-BN-ReLUスタックの統合などにより主にメモリ転送が軽減もしくは隠蔽されたためと考えられます。基本的にはTensorRTの実装はResNet50で使われるオペレータに関しては、ほぼハードウェアの性能を引き出しきっていると言えそうです。
また、show_summary_reportによって、オペレータ単位で集計した結果を見ることができます。
cch.show_summary_report(unit='M', mode='md', columns=c3.SummaryColumns.ALL)
この例では、デフォルトの項目に加えて、それぞれのオペレータがNN全体の理論計算量・理論メモリ転送量の何%を占めているかを報告しています。
| Layer type | # Layers | MFLOPs | MemRead MiB | MemWrite MiB | MemR+W MiB | FLOPs (%) | MemRead (%) | MemWrite (%) | MemR+W (%) |
|---|---|---|---|---|---|---|---|---|---|
| Convolution2DFunction | 53 | 3855.925 | 128.138 | 40.387 | 168.524 | 99.255% | 49.907% | 29.425% | 42.772% |
| FixedBatchNormalization | 53 | 10.587 | 40.589 | 40.387 | 80.976 | 0.273% | 15.809% | 29.425% | 20.552% |
| ReLU | 49 | 9.082 | 34.645 | 34.645 | 69.289 | 0.234% | 13.493% | 25.241% | 17.586% |
| MaxPooling2D | 1 | 1.606 | 3.062 | 0.766 | 3.828 | 0.041% | 1.193% | 0.558% | 0.972% |
| Add | 16 | 5.519 | 42.109 | 21.055 | 63.164 | 0.142% | 16.401% | 15.34% | 16.031% |
| AveragePooling2D | 1 | 0.1 | 0.383 | 0.008 | 0.391 | 0.003% | 0.149% | 0.006% | 0.099% |
| Reshape | 1 | 0.0 | 0.0 | 0.0 | 0.0 | 0.0% | 0.0% | 0.0% | 0.0% |
| LinearFunction | 1 | 2.049 | 7.824 | 0.004 | 7.828 | 0.053% | 3.047% | 0.003% | 1.987% |
| Softmax | 1 | 0.003 | 0.004 | 0.004 | 0.008 | 0.0% | 0.001% | 0.003% | 0.002% |
| total | 176 | 3884.872 | 256.754 | 137.254 | 394.008 | 100.0% | 100.0% | 100.0% | 100.0% |
ResNet50の場合は、総演算量の99%を、総メモリ転送量の43%をConvolutionが占めています。したがって、Convolutionに何かしらのアクセラレータを使うことができれば全体の性能向上に大きく寄与しそうです。しかし仮に強力な演算器によってConv演算だけをどれだけ速くできても、メモリ帯域が足を引っ張るため倍以上に速くはなりません。一方でBNとReLUだけで全体のメモリ転送量の40%近くを占めてるので、Conv-BN-ReLUを1つのCUDAカーネルにスタックできればこの40%分（からBNのβ/γのロード分を除いた分）をほぼまるっと減らせそうです。
このようなかなり具体的な議論が、理論計算量が求まるとできるようになります。
（もちろん、繰り返しですが厳密な議論は常に実環境上でのベンチマーク測定に基づいて行われるべきではあります。）
機能としては下記のようなものがあります
-
層ごと・オペレータの種類ごとのレポート（上述）
- csv, markdown, テキストでのテーブルで出力
- 補助単位(K,M,G,T)の指定、小数点以下桁数の指定
- 報告する項目の指定（全体に占める割合%、層のパラメータetc）
- 標準サポートされていない種類のオペレータの理論計算量算出ルールの外挿
- 層ごとの理論計算量レポートのデータへの直接アクセス
-
層ごとに、パラメータ(convでのk, padなど)やスタックトレースを保存
- report中の層が、ソースコード中のどの呼び出しに対応するか調べるのに便利です
- サポートされていないオペレータは単純に無視される
より詳細な使い方はGithubのREADMEをご参照ください。
また、オペレータごとのFLOPsおよびメモリ転送量の算出基準もドキュメント化されています。
開発と今後の展望
githubで開発していて、MITライセンスです。
今サポートされているオペレータとしては、ImageNet系(VGG, GoogLeNet, ResNet)、およびChainerCV(0.10.0)で使われているものを中心に整備しています。
しばらくはよく使われるようなオペレータを中心にサポート対象を増やす方向に開発が進むと思います。
それにあたって、内部的なインタフェースを中心に大きな変更が加わる可能性も高いです。
特にRNN系の対応が当面の課題です。
その他、あらゆる形での貢献を歓迎していますので、ぜひGithub上でディスカッションしましょう。
Special Thanks
本質的な仕組みは同僚のt-abeさんが考案したものです。
おわりに
Preferred Networksは深層学習技術の実世界応用に熱く取り組める人材を募集しています！
*1: FLOPS(FLoating-point Operation Per Second; 速度)ではないことに注意してください。FLoating-point OPerationSで、単にOPS(OPerationS)と呼ばれることもある単位で、単純に演算の回数を表す単位です。ならOPSでいいじゃんって話ですが、添字計算など直接出力に寄与しない演算を除くことを明示する意図もこめてFLOPsを使っています（HPCの文脈でのFLOPSも入力値から出力値を計算する本質的な命令の数だけを数えていてその他は含まれていないので）。