tf.summary の使用箇所を TF 2.0 に移行する

TensorFlow.org で表示 Google Colab で実行 GitHub でソースを表示 ノートブックをダウンロード

注意: このドキュメントは、TensorFlow 1.x TensorBoard に精通している方で大規模な TensorFlow コードベースを TensorFlow 1.x から 2.0 に移行したい方を対象としています。TensorBoard にまだ新しい方は、基礎ドキュメントをご覧ください。tf.keras を使用している場合は、TensorFlow 2.0 にアップグレードするための作業が必要ない場合があります。

import tensorflow as tf
2022-08-08 16:19:18.458555: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2022-08-08 16:19:19.410822: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvrtc.so.11.1: cannot open shared object file: No such file or directory
2022-08-08 16:19:19.411138: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvrtc.so.11.1: cannot open shared object file: No such file or directory
2022-08-08 16:19:19.411153: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.

TensorFlow 2.0 では、TensorBoard での視覚化に使用する要約データを書き込む際の tf.summary API が大幅に変更されています。

変更点

tf.summary API を 2 つのサブ API として考えると良いでしょう。

  • 個別の要約を記録するための一連の演算 - summary.scalar()summary.histogram()summary.image()summary.audio()、および summary.text()。これらはモデルコードからインラインで呼び出されます。
  • 上記の個別の要約を収集して特別にフォーマットされたログファイル(TensorBoard が読み取って視覚化を生成するファイル)に書き込む書き込みロジック。

TF 1.x の場合

上記の 2 つは、Session.run() でサマリー演算の出力をフェッチし、FileWriter.add_summary(output, step) で呼び出す、というように手動でつなぐ必要がありました。v1.summary.merge_all() 演算によって、グラフコレクションを使ってすべてのサマリー演算出力を集計するという方法で、この処理が簡単になっていますが、Eager execution と制御フローではあまりよく機能しなかったため、TF 2.0 には特に適していません。

TF 2.X の場合

上記の 2 つは密接に統合されており、個別の tf.summary 演算は実行時に直ちにデータを書き込むようになっています。モデルコードから API を使用する方法はあまり変わっていませんが、Eager execution との相性が改善されており、ほかのグラフモードとの互換性もそのままです。この 2 つの API を統合することで、summary.FileWriter は TensorFlow 実行コンテキストの一環となり、tf.summary 演算で直接アクセスできるため、ライターの構成が外見的な主な違いと言えます。

次は、TF 2.x のデフォルトのモードである Eager execution を使用した例です。

writer = tf.summary.create_file_writer("/tmp/mylogs/eager")

with writer.as_default():
  for step in range(100):
    # other model code would go here
    tf.summary.scalar("my_metric", 0.5, step=step)
    writer.flush()
ls /tmp/mylogs/eager
events.out.tfevents.1659975563.kokoro-gcp-ubuntu-prod-1801920270.9052.0.v2

次は、tf.function グラフ実行の使用例です。

writer = tf.summary.create_file_writer("/tmp/mylogs/tf_function")

@tf.function
def my_func(step):
  with writer.as_default():
    # other model code would go here
    tf.summary.scalar("my_metric", 0.5, step=step)

for step in tf.range(100, dtype=tf.int64):
  my_func(step)
  writer.flush()
ls /tmp/mylogs/tf_function
events.out.tfevents.1659975564.kokoro-gcp-ubuntu-prod-1801920270.9052.1.v2

次は、レガシー TF 1.x グラフ実行の使用例です。

g = tf.compat.v1.Graph()
with g.as_default():
  step = tf.Variable(0, dtype=tf.int64)
  step_update = step.assign_add(1)
  writer = tf.summary.create_file_writer("/tmp/mylogs/session")
  with writer.as_default():
    tf.summary.scalar("my_metric", 0.5, step=step)
  all_summary_ops = tf.compat.v1.summary.all_v2_summary_ops()
  writer_flush = writer.flush()


with tf.compat.v1.Session(graph=g) as sess:
  sess.run([writer.init(), step.initializer])

  for i in range(100):
    sess.run(all_summary_ops)
    sess.run(step_update)
    sess.run(writer_flush)
ls /tmp/mylogs/session
events.out.tfevents.1659975564.kokoro-gcp-ubuntu-prod-1801920270.9052.2.v2

コードを変換する

既存の tf.summary の使用箇所を TF 2.x API に変換する作業を確実に変換することは困難であるため、tf_upgrade_v2 スクリプトは、すべてを tf.compat.v1.summary に書き換えることだけを行い、自動的に TF 2.x の動作を有効化することはありません。

部分移行

tf.compat.v1.summary.scalar() といった TF 1.x サマリー API のロギング演算に大きく依存しているモデルコードを使用するユーザーが TF 2.x により簡単に移行できるようにするには、先にライター API のみを移行して、後でモデルコード内の個別の TF 1.x サマリー演算を完全に移行することができます。

このような移行をサポートするために、tf.compat.v1.summary は以下の条件で TF 2.x 相当に自動的に転送されます。

TF 2.x サマリー実装が呼び出されると、戻り値は空のバイト文字列テンソルになり、サマリーの書き込みが重複しないようにされることに注意してください。また、入力引数のフォワーディングはベストエフォートであり、すべての引数が保持されるとは限りません(たとえば、family 引数はサポートされていますが、collections は取り除かれます)。

以下は、tf.compat.v1.summary.scalartf.summary.scalar の動作を呼び出す例です。

# Enable eager execution.
tf.compat.v1.enable_v2_behavior()

# A default TF 2.x summary writer is available.
writer = tf.summary.create_file_writer("/tmp/mylogs/enable_v2_in_v1")
# A step is set for the writer.
with writer.as_default(step=0):
  # Below invokes `tf.summary.scalar`, and the return value is an empty bytestring.
  tf.compat.v1.summary.scalar('float', tf.constant(1.0), family="family")

完全移行

TF 2.x に全移行するには、コードを次のように適合させる必要があります。

  1. サマリー演算を使用するには、.as_default() によるデフォルトのライターセットが存在する必要がある

    • つまり、演算を Eager で実行するか、グラフ構造で演算を使用する
    • デフォルトのライターがない場合、サマリー演算はサイレントの no-op になる
    • デフォルトのライターは(まだ)@tf.function 実行境界に伝搬しません。関数がトレースされた場合にのみ検出されるため、関数の本文で writer.as_default() を呼び出し、@tf.function が使用される限りライターオブジェクトが存在し続けられるようにすることが、ベストプラクティスと言えます。
  2. 「step」値は step 引数で各演算に渡される必要がある

    • TensorBoard には、時系列としてデータをレンダリングするステップ値が必要です
    • TF 1.x のグローバルステップは削除されており、各演算は読み取るための希望する step 変数を知っておくために、明示的に渡す必要があります
    • ボイラープレートを減らすために、デフォルトステップを登録するための実験的サポートはtf.summary.experimental.set_step() として提供されていますが、これは暫定機能であり、予告なく変更される場合があります
  3. 個々のサマリー演算の関数シグネチャが変更されている

    • 戻り値はブール型になっています(要約が実際に書き込まれたかどうかを示す)
    • 2 番目のパラメータ名(使用される場合)がtensor から data に代わっています
    • collections パラメータが削除されています。collections は TF 1.x のみのパラメータです
    • family パラメータが削除されています。tf.name_scope() を使用してください
  4. [レガシーグラフモードのみ / セッション実行ユーザー]

    • 最初に v1.Session.run(writer.init()) でライターを初期化します

    • v1.summary.all_v2_summary_ops() を使用して、現在のグラフに関するすべての TF 2.0 サマリー演算を取得します(Session.run() で実行するためなど)。

    • v1.Session.run(writer.flush()) でライターをフラッシュし、close() でも同様にフラッシュします

TF 1.x コードで tf.contrib.summary API を使用していた場合は、TF 2.0 API にはるかに似ているため、tf_upgrade_v2 スクリプトを使って、ほとんどの移行ステップ(および完全に移行できない使用箇所の発行警告またはエラー)を自動化できます。ほとんどにおいて、API 呼び出しを tf.compat.v2.summary に書き直すだけであるため、TF 2.0+ との互換性のみが必要である場合は、compat.v2 を削除して、tf.summary として参照するようにすることができます。

その他のヒント

上記の重要な分野に加え、一部の補助的な側面も変更されています。

  • 条件付き記録(「100 ステップごとにログ」など)が新しくなりました

    • 演算と関連するコードを制御するには、通常の if ステートメント(Eager モードと自動グラフ作成経由の @tf.function で機能)か、tf.cond でラップします
    • 要約のみを制御するには、新しい tf.summary.record_if() コンテキストマネージャを使用して、選択したブール条件を渡します
    • TF 1.x パターンを置き換えます

      if condition:
        writer.add_summary()
      
  • tf.compat.v1.Graph を直接書き込めない代わりにトレース関数を使用します

  • tf.summary.FileWriterCache の使用により、logdir ごとのグローバルライターのキャッシュが不要になります

    • ユーザーは、ライターオブジェクトの独自のキャッシュ/共有を実装するか、個別のライターを使用する必要があります(TensorBoard の後者のサポートは開発中です)
  • イベントファイルのバイナリ表現が変更されました

    • TensorBoard 1.x は、新しい形式をサポート済みです。この違いは、要約データをイベントファイルから手動で解析しているユーザーのみに影響します)
    • 要約データはテンソルバイトとして保存されるようになりました。tf.make_ndarray(event.summary.value[0].tensor) を使って numpy に変換できます。