こんにちは。TensorFLow Advent Calendar 2018 2日目の記事です。
今年の9月終わり頃、TPUv2をColaboratory上でTPU(v2)を利用できるようになりました。それに伴って、いくつかの記事でGPUとTPUの比較もされました。TPU初出の論文では、K80の15倍〜30倍とうたっていた1こともあり、実際のTPUの性能はどうなんだ、ということで注目を浴びましたが、みんな苦労しているようです。
知り合いのGooglerに聞いても、やはりいろいろチューニングのポイントがあるようです。実際、公式のサンプル(KerasでFashion MNIST)をもとにコードを書いてみたのですが、プロファイラを見てみると、、、
上図のUtilization of TPU Matrix Units を見ていただければわかるように、なにも工夫しない状態では Matrix Units のたった 4.1% しか使っていないことがわかります。
というわけで今回は、TPUを使う時に気をつけることがまとまっている公式のページ
の内容をざっくりとまとめたいと思います。プロファイラの使い方については、
を御覧ください。
Cloud TPU Performance Guid より、気をつけるべきこと
- TPUを構成する 128x128 の Matrix Unit を最大限に利用するために、バッチサイズ、特徴数、画像の大きさ等に気をつける
- TPUを使う場合、データのある次元が8の倍数になるように、さらにもう一つの次元が128の倍数になるように、自動的にpaddingされます。paddingされた分は単純に無駄な計算となってしまうため、バッチサイズや特徴量を128の倍数や8の倍数とします。
- どの次元が8の倍数で、どの次元が128の倍数となるかは、演算ごとに違います。
- XLAの特性を知っておく
- TPUを使う場合、計算グラフは一旦XLAと呼ばれるツールで最適化され、TPU用のコードに変換されるため、XLAの特性を知っておくと、最適化に有利です
- 1.のpaddingも、正確にはXLAの特性です。
- ブロードキャスティングに注意
- これはTPU特有の話ではありませんが、例えば ベクトルと行列 の和を計算する時、ベクトルは（同じ値を持つ）行列に変換されてから、行列と和を取られることになります。やりたい計算によっては、計算の順序を変えるだけで、ブロードキャストをしないですみます。
- TensorFlowの各演算子毎の特性を知っておく
- paddingされている次元方向での slice は極力避ける
- max_pool よりは avg_pool のほうが、勾配計算が速いので、代替可能な場合はavg_poolを使う
- tf.transpose はいつでも計算量がほぼゼロというわけではないので注意
- 対象の行列の性質（畳み込みネットワークのカーネルなのか、アクティベーションなのか etc） によって挙動が違う
- reshapeも同じくいつでも計算量がほぼゼロというわけではない
- 乱数については、uniform/Bernoulliはとても速い。Normalはちょっと速い。Categorical/Multinomialは遅い。
- tf.nn.batch_normalization ではなくて fused_batch_norm を使う
- tf.nn.depthwise_conv2d, tf.nn.separable_conv2d はまだ最適化されてないから使わないでくれ！
データの入力パイプラインについて注意すべきこと
行列演算が高速だと、データの入力がボトルネックとなることがあります。実際、最近公開された ImageNetを2.2分で学習しきった論文でも、入力パイプラインについて言及されています。
- tf.data API を使う
- 入力パイプラインの最後に
prefetchを入れておいて、入力画像の変換とネットワークの訓練が並列して実行できるようにしておく -
mapを使うときは、num_parallel_callsをCPU数を指定しておく - 前処理してからバッチを作る場合は、
map_and_batchの利用を検討する。とくにバッチサイズが大きい場合 - データがリモートにある場合は、
parallel_interleaveを使う - 簡単なユーザ定義関数は、mapを使ってベクトル化しておく
- データが小さい場合は
cacheを使う - 前処理でデータ量が増えるような場合は、
interleave、prefetch、shuffleは早い段階で実施しておく -
shuffleを使うのはrepeatの前。もしくは、shuffle_and_repeatを使う
TPUのプロファイリングツールを使う
いくらTPUの力を発揮しようとしても、プロファイリングができないと、真っ暗闇を懐中電灯なしで歩くようなものです。
TPUのプロファイリングには、Google謹製の Cloud TPU Profiler を利用します。Cloud TPU Profilerによって結果が収集され、TensorBoardを使って見ることができます。
Colaboratory上で作業をしている場合、例えばGCS上にログをはいておいてリモートから確認することもできますが、以下の様にngrokを使うことで、ローカルに溜まったログをTensorBoardで確認することもできて便利です2。
#@title TensorBoardをColaboratoryで利用する
# ngrokをダウンロード
get_ipython().system_raw(f"""
wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip
unzip -o ngrok-stable-linux-amd64.zip
""")
# TensorBoard と ngrok の起動
get_ipython().system_raw(f"""
tensorboard --logdir {LOG_DIR} --host 0.0.0.0 --port {TFBOARD_PORT} &
./ngrok http 6006 &
""")
! curl -s http://localhost:4040/api/tunnels | python3 -c \
"import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])"
実際にプロファイリングをするのは、capture_tpu_profileコマンドです。このコマンドは初期状態ではColaboratoryに含まれていませんが、pipでインストールすることができます。
pip install --upgrade "cloud-tpu-profiler>=1.12"
capture_tpu_profileコマンドを、TPUを使っている最中にバックグラウンドで動かすことで、プロファイリングしてくれます。Colaboratoryだと、バックグラウンドでのコマンド実行は多少厄介ですが、例えば以下の様にします。
get_ipython().system_raw(f"""
echo $( for i in `seq 10`; do
capture_tpu_profile --service_addr $COLAB_TPU_ADDR --logdir={LOG_DIR}
sleep 10
done ) &
""")
tpu_model.fit_generator(...)
上記のコマンドを実行したあとで、おもむろにTensorBoardを開くと、以下のようにオーバービューを見ることができます。
本来であれば、Step-Time Graph が描画され、Recommendation for Next Step に、チェックすべき項目が表示されるのですが、Kerasの場合は TensorFlowの global_step を使っていないためか、Step-Time Graph が表示されず、Recommendation for Next Step もエラーとなるようです。
とはいえ、Top 10... の箇所をみれば、どの演算に全体のどれくらいの時間がかかっているかもわかりますし、Toolsから input_pipeline_analyzerやtrace_viewerを使ってボトルネックを見つけ出すことができます。
input_pipeline_analyer
trace_viewer
まとめ
なんとかしてTPUを使いこなしましょう。