TensorFlow Profilerで機械学習アプリケーションをプロファイリングしパフォーマンス改善のヒントをもらってみた

データアナリティクス事業本部 機械学習チームの鈴木です。

TensorFlow Profilerを使ってモデル訓練のパフォーマンス改善をしたいことがあったので、改めて使い方についてまとめてみました。

TensorFlow Profilerとは

TensorFlowコードの実行をプロファイリングし、機械学習アプリケーションのパフォーマンスを数値化して、より最適化するためのツールです。

以下のガイドに紹介があります。

TensorFlowコードを使った機械学習アプリケーションを開発する際に、使用するメソッドやオプションの使用状況によって期待した性能が出ないことがありますが、実装を修正しより効果的にマシンパワーを使うため、プロファイリングによる数値化とそれを受けた改善が非常に重要になります。

充実したツールが搭載されており、詳細については以下のガイドをご確認ください。

検証用のログの作成

TensorFlow Profilerによるパフォーマンス改善を行うため、検証用のログを作成していきます。

サンプルのノートブックについて

Google Colabで以下の処理を実装し、TensorFlow Profilerでパフォーマンスをプロファイリングしてみました。

作成したノートブックは、以下のようにGitHubレポジトリで公開しております。ここではポイントをご紹介します。

実装のポイント

まず、tensorboard-plugin-profileをインストールしておく必要があります。

!pip install -U tensorboard_plugin_profile

プロファイルを取得するため、以下のようにTensorBoardコールバックを作成しました。

# Create a TensorBoard callback
logs = "logs/" + datetime.now().strftime("%Y%m%d-%H%M%S")

tboard_callback = tf.keras.callbacks.TensorBoard(log_dir = logs,
                                                 histogram_freq = 1,
                                                 profile_batch = (10,20))

特にprofile_batchでプロファイリングするバッチを指定しました。ここで、サンプリングできる値もしくは範囲を指定しないと、TensorBoardでプロファイル結果が出ないので注意が必要です。

オプションの詳細はtf.keras.callbacks.TensorBoardのガイドをご確認ください。

コールバックは以下のように学習時にfitメソッドに設定しました。

model.fit(
  train_ds,
  validation_data=val_ds,
  epochs=3,
  callbacks = [tboard_callback]
)

今回使ったTensorFlow Profilerのバージョンでは、ID名のディレクトリ配下にログを移動させる必要があるようでした。この事象は『Profiler : No Profile Data was Found in Collab code example · Issue #602 · tensorflow/profiler』で議論されているのを参考にしました。

ログの移動は以下のように実装しています。

import glob
import shutil

for log_subdir in glob.glob('./logs/*'):
  for log_file in glob.glob(f'{log_subdir}/*/events.out.tfevents.*'):
    shutil.move(log_file, log_subdir)

以下のような状態であればプロファイル結果がTensorBoardから確認できました。

作成したログのバリエーションについて

TensorFlow Profilerでの改善が分かりやすいように、以下のバリエーションの実装でログを作成しました。

1. サンプルノートの実装からtrain_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)の実装を抜いたもの
2. サンプルノートの実装そのもの
3. サンプルノートの実装にos.environ["TF_GPU_THREAD_MODE"] = 'gpu_private'の実装を追加したもの

後にも記載しますが、No.1, 2の間では明らかな改善が見られました。No.2, 3の間では明らかな改善は見られませんでした。

プロファイル結果の確認

前提

TensorBoardをGoogle Colabで表示して確認しました。

%tensorboardマジックコマンドで表示するのが一番手っ取り早いですが、記事執筆時点ではこのように表示しようとすると、意図せず403エラーを返すようになる事象が起こっていました。

以下で紹介頂いているoutput.serve_kernel_port_as_windowを使う方法でTensorBoardにアクセスしました。

• [Colab] Tensorboard が表示できない
• Colabでウェブアプリを実行する

キャッシングとプリフェッチなしの例

TensorBoardを開き、プルダウンからPROFILEを選んでPROFILEの画面を表示しました。（本来はページ上部のバーにPROFILEのリンクが出ているような印象でしたが、今回はでなかったのでこの方法で開きました。）

このページだけでも様々な情報が掲載されていて非常に有用ですが、特にRecommendation for Next Stepの欄を見ると次に何をすると良いか分かりやすかったです。

Google Chromeのブラウザの翻訳機能で翻訳してみると以下のような表示でした。いくつかの改善点があり、特に入力パイプラインにボトルネックがあるようでした。

確認するべきツールの指定や、参考になりそうなガイドへのリンクもあり、参照しやすいです。

Toolsからツールを切り替えられます。

tf_data_bottleneck_analysisを見てみました。experimentalのステータスの機能です。

プリフェッチにボトルネックがあることが分かります。

Suggestionで提案してくれているページは以下になります。

tf_data_bottleneck_analysisとガイドの内容を踏まえて、キャッシングとプリフェッチを使用すると速度の改善が見込めそうです。これは鉄板の工夫のようですね。

改善の実施例

以下のようにキャッシングとプリフェッチの設定を追加しました。

AUTOTUNE = tf.data.AUTOTUNE

train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

以下のようにエポック2以降の時間が短縮されました。

工夫実施前

工夫実施後

再度TensorFlow Profilerを開き、分析内容を元に、表示されているTF_GPU_THREAD_MODE環境変数の設定をコードに追加してみましたが、こちらはあまり効果はありませんでした。今回の検証環境で実現できる速度としてはこれくらいが限界なのかもしれません。

os.environ["TF_GPU_THREAD_MODE"] = 'gpu_private'

最後に

TensorFlow Profilerによる機械学習アプリケーションのパフォーマンスを数値化と改善についてご紹介しました。参考になりましたら幸いです。

ほかに参考にした資料

• TensorFlow Profilerを用いた学習時間の短縮 - ユニファ開発者ブログ
• Tensorflow Profilerを使って前処理 (tf.data.Dataset API) のパフォーマンスを改善する - メモ帳