はじめに
TPU環境のTensorFlow/Kerasのfit時間を短縮できる、'experimental_steps_per_execution'についての調査報告をする。
環境
- Google Colab TPU
- TensorFlow 2.3.0
- tf.keras 2.4.3
記事作成後にColabのTensorFlowが2.4.0になったが、'steps_per_execution'の名前で正式採用となり、'experimental_'なしで記述できるようになった。
experimental_steps_per_executionとはなにか
TensorFlow2.3.0のReleaseNoteによると、下記の通り。
Added experimental_steps_per_execution
arg to model.compile to indicate the number of batches to run per tf.function call. This can speed up Keras Models on TPUs up to 3x.
tf.kerasのmodel.compileに引数指定することで、一回のcallで実行するバッチの個数を指定できるもので、TPUでは最高３倍速くなる、とある。
TensorFlowのTPU使用ガイドによると、
To reduce python overhead, and maximize the performance of your TPU, try out the experimental experimental_steps_per_execution argument to Model.compile. Here it increases throughput by about 50%:
pythonのオーバーヘッドを減らして高速化するようだ。ここでは50%高速化される例がかいてある。
GPUでは高速化の効果が無いが、のちに述べる制限&副作用は適用される。つまり「機能として無効ではないがメリットがない」ということになる。
筆者がこの機能を知ったのは、TPUの謎のWanrningの解決法を探していた際に、これを設定すればよい、とどこかに書いてあったため。
当該のエラーはこちら。同様のWarningで悩んでいる方は、設定することをお勧めする。
WARNING:tensorflow:Callbacks method
on_train_batch_endis slow compared to the batch time
制限&副作用
- 入力データは最低でも'steps_per_epoch * epochs'でなければならない
- ちなみに切りのいい数値にしなくても切り捨てられることは無いようだ。
- tf.keras.callbacks.Callbackのon_batch_begin/on_batch_endの呼び出し回数が減る
- 通常はバッチ処理毎にこれらの関数が呼び出されるが、experimental_steps_per_executionを設定すると、その設定回数を処理し終わった時点で呼び出されるので結果的に回数が減る。
- fit時に最終的に表示される処理時間が不当に短くなることがある
- 必ずではないので何らかの条件があるようだ
- verbose=1の時に情報が更新される頻度が減る
効果確認
実験条件
- Google ColabのTPUで実施
- TPUでは初回のfitに異常に時間がかかるので、初回は捨て２回目のfitの時間を測定する。
- データはCIFAR10(32,32,3)
- ネットワークはパラメータ数の違うCNNを２つ
- High Model Total params: 10,165,578
- Low Model Total params: 2,187,434
- fitの設定
- steps_per_epoch = 1000
- batch_size = 256 と 32
- verbose = 0
- validationなし
experimental_steps_per_executionを6種類選んで実験している(0は無指定を意味する)。
verbose=1にしても同じような結果なのだが、無指定でのWarning表示のために正確な比較ができない可能があるので、verbose=0で実験している。
実験結果
batch_size=32の結果はこちら。
batch_size=256の結果はこちら。
全ての条件で500にした場合が一番結果が良かった。こちらの数値と無指定を基準とした比率を載せておく。
| batch_size | parameters | duration(sec) | % |
|---|---|---|---|
| 32 | 2M | 13.5 | 91.2 |
| 32 | 10M | 17.7 | 84.2 |
| 256 | 2M | 18.0 | 83.7 |
| 256 | 10M | 33.9 | 94.1 |
微妙な結果だが、高速化できることは事実である。
この実験では20とした場合には、逆に時間がかかってしまった。小さな数値を設定することは避けたほうがよさそうだ。
ちなみに前述のTensorFlowのガイドどおりの学習では11秒が6秒になって50%程度の高速化は嘘ではない。
この辺は条件次第のようだ。