ニューラルネットワークの演算量を計測する

Idein エンジニアの打田です。

今回はパフォーマンスカウンタを利用してニューラルネットワークの演算量を計測してみたので方法を共有したいと思います。

昨今は大規模なネットワークを学習するために GPU クラスタが利用されはじめていたり、推論についても GPU やディープニューラルネットワーク用のプロセッサなどの開発であったりと、ニューラルネットワークの演算負荷を効率よく処理するニーズは高まっています。

さてしかし、ニューラルネットワークの演算負荷というのはどのくらいのものなのでしょうか、あるプロセッサが利用できるときにどのニューラルネットワークがどれくらいで実行できるのか、あるいは、ある学習済みニューラルネットワークがあるときにどのプロセッサならどのくらいの時間で実行できるのか、プロセッサの性能とモデルの規模の関係に、おおよそのあたりがつけられるとニューラルネットワークの演算負荷というものについて体感できるようになってうれしい気がします。

FLOPS (FLoating point number Operations Per Seconds) はプロセッサの重要な性能指標ですし、演算負荷を評価するのに浮動小数点演算回数というのはよい指標でしょう(ニューラルネットワークに限ったことでありませんが)。

演算回数を評価する方法としては、各レイヤーの演算、入力サイズ、出力サイズの定義から想定される演算回数を見積っていく方法がまず思いつくかと思います。 実際 Netscope CNN Analyzer の MACCs (Multiply-ACCumulations) 計算や apaszke/torch-opCountertf.profiler.ProfileOptionBuilder.float_operation などはそのような方法で演算回数を評価しています。

今回は上記のような方法とは別の、ハードウェアに備わるパフォーマンスカウンタを利用して演算回数を計測する方法を紹介します。

この方法は実際の実行結果に基づく評価なので、利用するハードウェア、ソフトウェアの実装、アルゴリズム等の影響を強く受けるという面はありますが、利用しているディープラーニングフレームワークに関係なく計測できる、内部のネットワーク構成を考えることなく動く実行ファイルさえあれば計測できる、実際の計測結果が得られるため見積りの見落しがないかの確認にも利用できる、といった利点があります。

今回は chainer に含まれている ImageNet 用のネットワークの演算回数をいくつか評価してみようと思います。

計測に利用したコードは https://github.com/Idein/dnn-flop-count に置いてあります。

まずは CPU での計測方法から紹介します。

perf: Linux profiling with performance counters でとるのが割と手軽です。 ただ CPU によって使えるカウンタが違うという点には注意が必要です。例えば Intel Haswell 世代の CPU ではそもそもこの目的で利用できるカウンタがありません Counting Floating Point Operations on Intel Haswell - PAPITopics:SandyFlops - PAPIDocs

Skylake/Kabylake であれば単精度浮動小数点演算命令の数は以下のイベントで取得できるので、今回はたまたま社内にあった Skylake 世代の CPU を利用して測定しました。

  • FP_ARITH:SCALAR_SINGLE
    • umask: 0x02
    • event: 0xc7
  • FP_ARITH:128B_PACKED_SINGLE
    • umask: 0x08
    • event: 0xc7
  • FP_ARITH:256B_PACKED_SINGLE
    • umask: 0x20
    • code: 0xc7
  • FP_ARITH:512B_PACKED_SINGLE
    • umask: 0x80
    • code: 0xc7

イベントの code と umask は Intel® 64 and IA-32 Architectures Developer’s Manual: Vol. 3B から確認できます(お手持ちの CPU でよさそうなカウンタがないか眺めてみるのもいいですね)。

GoogLeNet の 推論の場合どうなるか確認してみましょう。 (今回は predict なので推論の見積りですが訓練の見積りをしたい場合は損失関数の呼出しをした後に backward するまでのコードで評価すればいいでしょう。)

$ perf stat -e r02c7,r08c7,r20c7,r80c7 python3 eval_imagenet.py googlenet 

 Performance counter stats for 'python3 eval_imagenet.py googlenet':

          85770094      r02c7                                                       
          47396811      r08c7                                                       
         372309776      r20c7                                                       
                 0      r80c7                                                       

       1.598643946 seconds time elapsed

PACK 分を考慮すると FP_ARITH:128B_PACKED_SINGLE は4演算相当 FP_ARITH:256B_PACKED_SINGLE は8演算 FP_ARITH:512B_PACKED_SINGLE は16演算相当ですから perf-x オプションと awk で集計してみましょう。

$ perf stat -e r02c7,r08c7,r20c7,r80c7 -x, python3 eval_imagenet.py googlenet 2>&1 | awk -F, '/r02c7/ {sum+=$1} /r08c7/ {sum+=4*$1} /r20c7/ {sum+=8*$1} /r80c7/ {sum+=16*$1} END {print sum}'
3254293004

1回の推論で3G超、30億回以上の演算ですね。この演算数を捌くにはやはりハイパワーなプロセッサが必要ですね、そう Raspberry Pi に載っている GPU Video Core IV のような。

perf を使うと実行ファイルさえあれば計測できて便利ですが、ネットワーク内のレイヤー毎の演算回数や各関数毎の演算回数が得られると演算のボトルネックなどもわかって便利ですね。

perfmon2libpfm4 には Python binding があるため、これを Python から呼び出して使うと値の取得が楽になります。

chainer を使う場合 chainer.FunctionHookchainer.Function の呼出しを hook できるので、それを利用しましょう。 https://github.com/Idein/dnn-flop-count/blob/master/perf_counter.py にある CounterHook を使うことで以下のようなコードで実行時の各関数の演算回数を把握できます。

with chainer.using_config('train', False):
    with chainer.using_config('enable_backprop', False):
        with CounterHook() as counter:
            model.predict(image, oversample=False)
for fn, float_ops in counter.call_history:
    print('"{}","{}"'.format(fn, float_ops))

実行例

$ python eval_imagenet.py --count-by functions googlenet
"Convolution2DFunction","30507008"
"ReLU","0"
"MaxPooling2D","0"
"LocalResponseNormalization","18919576"
...snip...
"Convolution2DFunction","1315552"
"Concat","0"
"ReLU","0"
"AveragePooling2D","51200"
"Reshape","0"
"LinearFunction","258750"
"Softmax","5995"

レイヤー単位で数える場合はレイヤーに hook する機能はないため、少し行儀の悪い方法にはなりました、レイヤー毎に回数を調べることもできます。 このあたりは chainer の computational_graph 機能と組み合わせるともっと格好いい使い方ができるかもしれません。

NVIDIA GPU で計測したい場合には nvprof が perf のように利用できます。

9. Metrics Reference - Profiler :: CUDA Toolkit Documentation を見ると flop_count_sp がまさに単精度浮動小数点演算回数を求めるのに利用できます(FMA 演算が2命令分でカウントされていますが、CPU での計測も FMA を2命令分でカウントしているため今回の目的には合致します)。

そのままだと CUDA kernel 毎に、呼び出し回数(3列目)と平均(8列目)が出力されるので nvprof で csv 出力させつつ gawk で平均と呼出し回数を合算します。

$ nvprof --csv --metrics flop_count_sp python3 eval_imagenet.py resnet50 --gpu=0 2>&1 | gawk -v 'FPAT=([^,]+)|(\"[^"]+\")' '{sum+=$3*$8} END {print sum}'
2994793235

nvprof 便利ですね。

nvprof で取得できる情報は CUPTI (CUDA Performance Tools Interface) の API で取得可能できるようなので、うまくラッパーを書けばレイヤーや関数毎の演算回数を出力することもできそうです。

ニューラルネットワークの訓練で発生する熱が気になる季節になってきましたが、熱の原因であるところの演算量に思いを馳せてみるのもいいかもしれません。