TensorBoard のスカラー: Keras でトレーニングメトリックをログする

TensorFlow.org で表示 Google Colab で実行 GitHub でソースを表示

概要

機械学習には、損失などの主要なメトリック、そしてトレーニングが進むにつれそれらがどのように変化するかを理解することが必ず伴います。こういったメトリックは、過適合が発生していないか、不要に長くトレーニングしていないかなどを理解する上で役立ちます。そのため、モデルのデバッグや改善を行いやすくするために、トレーニングの実行から得られるメトリックを比較すると良いでしょう。

TensorBoard の Time Series ダッシュボードでは、単純な API を使用して簡単にこれらのメトリックを可視化できます。このチュートリアルでは非常に基本的な例を使用して、Keras モデルを開発する際に API と TensorBoard を使用する方法を説明します。Keras TensorBoard コールバックと TensorFlow Summary API を使用して、デフォルトとカスタムのスカラーを可視化する方法を学習します。

セットアップ

# Load the TensorBoard notebook extension.
%load_ext tensorboard

from datetime import datetime
from packaging import version

import tensorflow as tf
from tensorflow import keras
from keras import backend as K

import numpy as np

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

TensorFlow version:  2.8.2

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

簡単な回帰用データをセットアップする

これから Keras を使用して回帰を計算し、データセットペアにおける最適な適合線を見つけたいと思います。（この種の問題でニューラルネットワークと勾配降下を使用するのはやり過ぎではありますが、例を理解しやすくなります。）

TensorBoard を使用して、エポック間でトレーニングとテストの損失がどのように変化するのかを観測します。時間の経過とともにトレーニングとテストの損失が下降して安定化する様子を確認してください。

まず、おおまかに線 y = 0.5x + 2 に沿って 1000 個のデータポイントを生成し、それらをトレーニングセットとテストセットに分割します。ここでは、ニューラルネットワークがこの関係を学習することを目標とします。

data_size = 1000
# 80% of the data is for training.
train_pct = 0.8

train_size = int(data_size * train_pct)

# Create some input data between -1 and 1 and randomize it.
x = np.linspace(-1, 1, data_size)
np.random.shuffle(x)

# Generate the output data.
# y = 0.5x + 2 + noise
y = 0.5 * x + 2 + np.random.normal(0, 0.05, (data_size, ))

# Split into test and train pairs.
x_train, y_train = x[:train_size], y[:train_size]
x_test, y_test = x[train_size:], y[train_size:]

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

モデルの定義、トレーニング、および評価の準備が整いました。

トレーニングしながら損失スカラーをログするには、次のように行います。

  1. Keras TensorBoard コールバックを作成します。
  2. ログディレクトリを指定します。
  3. TensorBoard コールバックを Keras の Model.fit() に渡します。

TensorBoard は、ログディレクトリ階層からログデータを読み取ります。このノートブックでは、ルートログディレクトリは logs/scalars で、その後にタイムスタンプ付きのサブディレクトリが続きます。タイムスタンプがサブディレクトリに付加されることで、TensorBoard を使用してモデルでイテレーションを行う際に、トレーニングセットを簡単に識別して選択することができます。 

logdir = "logs/scalars/" + datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)

model = keras.models.Sequential([
    keras.layers.Dense(16, input_dim=1),
    keras.layers.Dense(1),
])

model.compile(
    loss='mse', # keras.losses.mean_squared_error
    optimizer=keras.optimizers.SGD(learning_rate=0.2),
)

print("Training ... With default parameters, this takes less than 10 seconds.")
training_history = model.fit(
    x_train, # input
    y_train, # output
    batch_size=train_size,
    verbose=0, # Suppress chatty output; use Tensorboard instead
    epochs=100,
    validation_data=(x_test, y_test),
    callbacks=[tensorboard_callback],
)

print("Average test loss: ", np.average(training_history.history['loss']))

Training ... With default parameters, this takes less than 10 seconds.
Average test loss:  0.042797307365108284

TensorBoard を使って損失を調べる

では、上記で使用したルートログディレクトリを指定して、TensorBoard を起動しましょう。

TensorBoard の UI が読み込まれるまで数秒ほどかかります。 

%tensorboard --logdir logs/scalars

TensorBoard に「No dashboards are active for the current data set（現在のデータセットにはアクティブなダッシュボードはありません）」というメッセージが表示されることがあります。これは、最初のログデータがまだ保存されていないためです。トレーニングが進むにつれ、Keras モデルはデータをログし始めるため、TensorBoard の再読み込みが定期的に行われ、スカラーメトリックが表示されるようになります。それまで待てない方は、右上にある再読み込み矢印をタップしてください。

トレーニングが進むにつれ、トレーニングと検証の損失が急速に下降し、安定する様子が見られます。実際、トレーニングが 25 エポックを終了したところであまり改善が見られなくなるため、その時点でトレーニングを停止することができます。

グラフをホバーし、特定のデータポイントを確認します。マウスを使用して拡大したり、部分的に選択して詳細を確認することも可能です。

左側に「Runs（実行）」セレクタがあります。1 つの「実行」は、1 ラウンドのトレーニングから得たログ一式を指します。このケースでは、これは Model.fit() の結果です。通常、長期間にわたってモデルの実験と開発が行われるため、実行の数は非常に多くなります。

Runs セレクタを使用して特定の実行を選択するか、トレーニングまたは検証のみから選択します。実行を比較すると、どのバージョンのコードで問題の解決が最適に行われているかを評価しやすくなります。

TensorBoard の損失グラフには、トレーニングと検証の両方で損失が一定して下降し、安定したことが示されています。つまり、モデルのメトリックは非常に良質である可能性が高いということです。では、このモデルが実際のデータでどのように動作するかを見てみましょう。

入力データ (60, 25, 2) がある場合、線 y = 0.5x + 2 は (32, 14.5, 3) となるはずです。モデルはこれと合致するでしょうか。

print(model.predict([60, 25, 2]))
# True values to compare predictions against: 
# [[32.0]
#  [14.5]
#  [ 3.0]]

[[32.148884 ]
 [14.562463 ]
 [ 3.0056725]]

よくできました！

カスタムスカラーをログする

動的学習率などのカスタムの値をログする場合はどうでしょうか。この場合は、TensorFlow Summary API を使用する必要があります。

回帰モデルを維持したまま、カスタム学習率をログします。次のように行ってください。

  1. tf.summary.create_file_writer() を使ってファイルライターを作成します。
  2. カスタム学習率の関数を定義します。この関数は、Keras の LearningRateScheduler コールバックに渡されます。
  3. 学習率関数内に、カスタム学習率をログするための tf.summary.scalar() を使用します。
  4. LearningRateScheduler コールバックを Model.fit() に渡します。

一般的に、カスタムスカラーをログするには、tf.summary.scalar() をファイルライターとともに使用する必要があります。この実行のデータを特定のディレクトリに書き込むのはファイルライターであり、tf.summary.scalar() を使用する際にファイルライターが暗黙的に使用されるためです。

logdir = "logs/scalars/" + datetime.now().strftime("%Y%m%d-%H%M%S")
file_writer = tf.summary.create_file_writer(logdir + "/metrics")
file_writer.set_as_default()

def lr_schedule(epoch):
  """
  Returns a custom learning rate that decreases as epochs progress.
  """
  learning_rate = 0.2
  if epoch > 10:
    learning_rate = 0.02
  if epoch > 20:
    learning_rate = 0.01
  if epoch > 50:
    learning_rate = 0.005

  tf.summary.scalar('learning rate', data=learning_rate, step=epoch)
  return learning_rate

lr_callback = keras.callbacks.LearningRateScheduler(lr_schedule)
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)

model = keras.models.Sequential([
    keras.layers.Dense(16, input_dim=1),
    keras.layers.Dense(1),
])

model.compile(
    loss='mse', # keras.losses.mean_squared_error
    optimizer=keras.optimizers.SGD(),
)

training_history = model.fit(
    x_train, # input
    y_train, # output
    batch_size=train_size,
    verbose=0, # Suppress chatty output; use Tensorboard instead
    epochs=100,
    validation_data=(x_test, y_test),
    callbacks=[tensorboard_callback, lr_callback],
)

もう一度 TensorBoard を確認しましょう。

%tensorboard --logdir logs/scalars

左側の「Runs」セレクタを使用すると、<timestamp>/metrics 実行が 1 つあります。この実行を選択すると「learning rate（学習率）」グラフが表示され、この実行における学習率の進行を確認することができます。

また、この実行のトレーニングと検証の損失曲線を以前の実行と比較することもできます。また、エポックによっては、学習率スケジュールが離散値を返すことがありますが、学習率プロットは滑らかに見える場合があることに気付くかもしれません。TensorBoard には平滑化パラメーターがあり、平滑化されていない値を表示するには、ゼロまで下げる必要がある場合があります。

このモデルの出来栄えはどうでしょうか。

print(model.predict([60, 25, 2]))
# True values to compare predictions against: 
# [[32.0]
#  [14.5]
#  [ 3.0]]

[[31.958094 ]
 [14.482997 ]
 [ 2.9993598]]

バッチレベルのロギング

まず MNIST データセットを読み込んでデータを正規化し、画像を 10 個のクラスに分類する単純な Keras モデルを作成する関数を記述しましょう。

mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

def create_model():
  return tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(512, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
  ])

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
11501568/11490434 [==============================] - 0s 0us/step

バッチレベルの瞬時ロギング

バッチレベルでメトリクスを瞬時にロギングすると、エポックごとにトレーニング中のバッチ間の変動レベルを見ることができます。これは、デバッグの際に役立ちます。

サマリーライターを別のログディレクトリにセットアップします。

log_dir = 'logs/batch_level/' + datetime.now().strftime("%Y%m%d-%H%M%S") + '/train'
train_writer = tf.summary.create_file_writer(log_dir)

バッチレベルのロギングを有効にするには、モデルのクラス定義で train_step() をオーバーライドしてカスタム tf.summary メトリクスを定義し、サマリーライターのコンテキストに含める必要があります。これは、以下に示すように、サブクラス化されたモデル定義に単純に結合するか、以前の関数 API モデルを編集して拡張できます。

class MyModel(tf.keras.Model):
  def __init__(self, model):
    super().__init__()
    self.model = model

  def train_step(self, data):
    x, y = data
    with tf.GradientTape() as tape:
      y_pred = self.model(x, training=True)
      loss = self.compiled_loss(y, y_pred)
      mse = tf.keras.losses.mean_squared_error(y, K.max(y_pred, axis=-1))
    self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
    with train_writer.as_default(step=self._train_counter):
      tf.summary.scalar('batch_loss', loss)
      tf.summary.scalar('batch_mse', mse)
    return self.compute_metrics(x, y, y_pred, None)

  def call(self, x):
    x = self.model(x)
    return x

# Adds custom batch-level metrics to our previous Functional API model
model = MyModel(create_model())
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

エポックレベルとバッチレベルのメトリクスをログディレクトリにログし、選択した batch_size で model.fit() を呼び出すように TensorBoard コールバックを定義します。 

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir)

model.fit(x=x_train, 
          y=y_train,
          epochs=5,
          batch_size=500, 
          validation_data=(x_test, y_test), 
          callbacks=[tensorboard_callback])

Epoch 1/5
120/120 [==============================] - 5s 36ms/step - loss: 0.4379 - accuracy: 0.8788 - val_loss: 0.2041 - val_accuracy: 0.9430
Epoch 2/5
120/120 [==============================] - 4s 31ms/step - loss: 0.1875 - accuracy: 0.9471 - val_loss: 0.1462 - val_accuracy: 0.9591
Epoch 3/5
120/120 [==============================] - 3s 27ms/step - loss: 0.1355 - accuracy: 0.9613 - val_loss: 0.1170 - val_accuracy: 0.9670
Epoch 4/5
120/120 [==============================] - 3s 27ms/step - loss: 0.1058 - accuracy: 0.9694 - val_loss: 0.0954 - val_accuracy: 0.9723
Epoch 5/5
120/120 [==============================] - 3s 27ms/step - loss: 0.0872 - accuracy: 0.9752 - val_loss: 0.0843 - val_accuracy: 0.9749
<keras.callbacks.History at 0x7fce165a2fd0>

新しいログディレクトリで TensorBoard を開き、エポックレベルとバッチレベルのメトリクスを確認します。

%tensorboard --logdir logs/batch_level

バッチレベルの累積ロギング

バッチレベルのロギングは、累積して実装することも可能です。バッチのメトリクスと前のメトリクスを平均化することで、バッチレベルのメトリクスをロギングする際に、より滑らかなトレーニング曲線が描かれます。

サマリーライターを別のログディレクトリにセットアップします。

log_dir = 'logs/batch_avg/' + datetime.now().strftime("%Y%m%d-%H%M%S") + '/train'
train_writer = tf.summary.create_file_writer(log_dir)

バッチごとにログできるステートフルメトリクスを作成します。

batch_loss = tf.keras.metrics.Mean('batch_loss', dtype=tf.float32)
batch_accuracy = tf.keras.metrics.SparseCategoricalAccuracy('batch_accuracy')

前と同様に、オーバーライドされた train_step にカスタム tf.summary メトリクスを追加します。バッチレベルのロギングを累積して行うために、定義したステートフルメトリクスを使って、各トレーニングステップのデータに基づく累積結果を計算します。

class MyModel(tf.keras.Model):
  def __init__(self, model):
    super().__init__()
    self.model = model

  def train_step(self, data):
    x, y = data
    with tf.GradientTape() as tape:
      y_pred = self.model(x, training=True)
      loss = self.compiled_loss(y, y_pred)
    self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
    batch_loss(loss)
    batch_accuracy(y, y_pred)
    with train_writer.as_default(step=self._train_counter):
      tf.summary.scalar('batch_loss', batch_loss.result())
      tf.summary.scalar('batch_accuracy', batch_accuracy.result())
    return self.compute_metrics(x, y, y_pred, None)

  def call(self, x):
    x = self.model(x)
    return x

# Adds custom batch-level metrics to our previous Functional API model
model = MyModel(create_model())
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

前と同様に、TensorBoard コールバックを定義して、選択した batch_size で model.fit() を呼び出します。 

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir)

model.fit(x=x_train, 
          y=y_train,
          epochs=5,
          batch_size=500, 
          validation_data=(x_test, y_test), 
          callbacks=[tensorboard_callback])

Epoch 1/5
120/120 [==============================] - 4s 27ms/step - loss: 0.4266 - accuracy: 0.8813 - val_loss: 0.2055 - val_accuracy: 0.9415
Epoch 2/5
120/120 [==============================] - 3s 26ms/step - loss: 0.1864 - accuracy: 0.9476 - val_loss: 0.1417 - val_accuracy: 0.9613
Epoch 3/5
120/120 [==============================] - 3s 27ms/step - loss: 0.1352 - accuracy: 0.9614 - val_loss: 0.1148 - val_accuracy: 0.9665
Epoch 4/5
120/120 [==============================] - 3s 26ms/step - loss: 0.1066 - accuracy: 0.9702 - val_loss: 0.0932 - val_accuracy: 0.9716
Epoch 5/5
120/120 [==============================] - 3s 27ms/step - loss: 0.0859 - accuracy: 0.9749 - val_loss: 0.0844 - val_accuracy: 0.9754
<keras.callbacks.History at 0x7fce15c39f50>

新しいログディレクトリで TensorBoard を開き、エポックレベルとバッチレベルのメトリクスを確認します。

%tensorboard --logdir logs/batch_avg

これで完了です！様々なユースケースに使用するカスタムトレーニングメトリクスを TensorBoard に作成する方法を学習しました。