TensorFlow グラフを調査する

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

概要

TensorBoard の Graphs ダッシュボードは、TensorFlow モデルを調べるために使用できる強力なツールです。モデルの構造の概念的なグラフを素早く表示し、意図した設計と一致することを確認することができます。また、演算レベルのグラフも表示できるため、TensorFlow がどのようにプログラムを理解しているかを把握することができます。演算レベルグラフを調査すると、モデルの変更方法に関する洞察を得ることができます。たとえば、トレーニングの進行速度が思ったより遅い場合に、モデルを再設計することが可能となります。

このチュートリアルでは、グラフの診断データを生成して TensorBoard の Graph ダッシュボードで視覚化する方法を簡単に説明します。チュートリアルでは、Fashion-MNIST データセットの簡単な Keras Sequential モデルを定義してトレーニングし、モデルグラフのログと調査の方法を学習します。また、トレーニング API を使用して、新しい tf.function 注釈を使って作成した関数のグラフデータを生成します。

セットアップ

# Load the TensorBoard notebook extension.
%load_ext tensorboard

from datetime import datetime
from packaging import version

import tensorflow as tf
from tensorflow import keras

print("TensorFlow version: ", tf.__version__)
assert version.parse(tf.__version__).release[0] >= 2, \
    "This notebook requires TensorFlow 2.0 or above."

TensorFlow version:  2.2.0

import tensorboard
tensorboard.__version__

'2.2.1'

# Clear any logs from previous runs
rm -rf ./logs/

Keras モデルを定義する

この例では、単純な 4 つのレイヤーを使った Sequential モデルを分類器とします。

# Define the model.
model = keras.models.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(32, activation='relu'),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10, activation='softmax')
])

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'])

トレーニングデータをダウンロードして準備します。

(train_images, train_labels), _ = keras.datasets.fashion_mnist.load_data()
train_images = train_images / 255.0

モデルをトレーニングしてデータをログする

トレーニングの前に、ログディレクトリを指定して Keras TensorBoard コールバックを定義します。Model.fit() にこのコールバックを渡すことで、グラフデータが TensorBoard で視覚化できるようにログされるようにします。

# Define the Keras TensorBoard callback.
logdir="logs/fit/" + datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)

# Train the model.
model.fit(
    train_images,
    train_labels, 
    batch_size=64,
    epochs=5, 
    callbacks=[tensorboard_callback])

Epoch 1/5
938/938 [==============================] - 2s 2ms/step - loss: 0.6955 - accuracy: 0.7618
Epoch 2/5
938/938 [==============================] - 2s 2ms/step - loss: 0.4877 - accuracy: 0.8296
Epoch 3/5
938/938 [==============================] - 2s 2ms/step - loss: 0.4458 - accuracy: 0.8414
Epoch 4/5
938/938 [==============================] - 2s 2ms/step - loss: 0.4246 - accuracy: 0.8476
Epoch 5/5
938/938 [==============================] - 2s 2ms/step - loss: 0.4117 - accuracy: 0.8508
<tensorflow.python.keras.callbacks.History at 0x7f656ecc3fd0>

演算レベルグラフ

TensorBoard を起動し、UI が読み込まれるまで数秒ほど待ってから、上部にある「Graph」をタップして、Graphs ダッシュボードを選択します。 

%tensorboard --logdir logs

オプションで TensorBoard.dev を使用して、ホストされた共有可能な実験を作成することもできます。 

!tensorboard dev upload \
  --logdir logs \
  --name "Sample op-level graph" \
  --one_shot

デフォルトでは、TensorBoard には演算レベルグラフが表示されます（左側に、「Default」タグが選択されているのがわかります）。グラフが反転していることに注意してください。データは下から上に向かって流れているため、コードと比較すると逆さまになっています。それでも、グラフが Keras モデルに緊密に一致しており、ほかの計算ノードへのエッジが追加されていることはわかるでしょう。

グラフは非常に大きいことが多いですが、グラフ表示を操作することができます。

  • スクロールで拡大縮小
  • ドラッグでパン
  • ダブルクリックでノード展開の切り替え（ノードはほかのノードのコンテナである場合があります）

また、ノードをクリックしてメタデータを表示することもできます。こうすると、入力、出力、およびその他の詳細を確認することができます。

概念グラフ

TensorBoard には、実行グラフのほかに概念グラフも表示されます。これは Keras モデルのみのビューです。保存したモデルを再利用する際に構造を調査または検証する場合に役立ちます。

概念グラフを表示するには、「keras」タグを選択します。折り畳まれた Sequential ノードが表示されますが、それをダブルクリックすると、モデルの構造を確認できます。


tf.function のグラフ

これまでの例では、Keras レイヤーを定義して Model.fit() を呼び出すことでグラフを作成するという、Keras モデルのグラフについて説明しました。

tf.function 注釈を使用して、「グラフの自動作成」、つまり Python 計算関数を高性能 TensorFlow グラフに変換しなければならない状況に遭遇することがあるかもしれません。こういった場合には、TensorFlow Summary Trace API を使用して、TensorBoard に視覚化するための自動グラフ作成関数をログすることができます。

Summary Trace API を使用するには、次を行います。

  • 関数を定義して、tf.function で注釈をつけます。
  • 関数呼び出しの直前に tf.summary.trace_on() を使用します。
  • profiler=True を渡して、グラフにプロファイル情報（メモリ、CPU 時間）を追加します。
  • サマリーファイルライターを使って、tf.summary.trace_export() を呼び出してログデータを保存します。

その後、TensorBoard を使用すると、関数の動作を確認することができます。

# The function to be traced.
@tf.function
def my_func(x, y):
  # A simple hand-rolled layer.
  return tf.nn.relu(tf.matmul(x, y))

# Set up logging.
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = 'logs/func/%s' % stamp
writer = tf.summary.create_file_writer(logdir)

# Sample data for your function.
x = tf.random.uniform((3, 3))
y = tf.random.uniform((3, 3))

# Bracket the function call with
# tf.summary.trace_on() and tf.summary.trace_export().
tf.summary.trace_on(graph=True, profiler=True)
# Call only one tf.function when tracing.
z = my_func(x, y)
with writer.as_default():
  tf.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)

%tensorboard --logdir logs/func

これで、TensorBoard が理解する通りの関数の構造を確認できるようになりました。「Profile」ラジオボタンをクリックすると、CPU とメモリの統計データを確認できます。