簡単 3分 TensorBoard in Google Colab (TensorFlow 2.x 使用)

TensorFlow を使用する際、学習の状態を可視化する TensorBoard がよく使われます。
Web 画面上で対話的に状態を見ることができ、とても便利なのですが Google Colaboratory 中から参照するには少し仕掛けが必要でした。
TensorFlow 2.x では Google Colab. 上で簡単に使用できるようになったようなので試してみます。

2020-12-06 追記　PyTorch でも TensorBoard は利用可能です。
PyTorch でも TensorBoard in Google Colab

実行環境

Google Colaboratory を使用します。

サンプルコード

Google Colab からの TensorBoard の利用については、 TensorFlow の公式ページに解説があります。
https://www.tensorflow.org/tensorboard/tensorboard_in_notebooks

TensorFlow 2.x の使用のマジックコマンドを実行します。

from __future__ import absolute_import, division, print_function, unicode_literals

try:
  # %tensorflow_version only exists in Colab.
  %tensorflow_version 2.x
except Exception:
  pass

次に TensorBoard の読み込みのためのマジックコマンドを実行します。

# Load the TensorBoard notebook extension
%load_ext tensorboard

MNIST（サンプルとしてよく使われる数字の画像データ）を使った簡単なモデルを作成します。keras だと非常にシンプルです。

import tensorflow as tf
from tensorflow.keras.callbacks import TensorBoard

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

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

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

TensorBoard のためのコールバック関数を model.fit() に与えます。
この辺りは通常の TensorBoard の使い方と同じです。

tf_callback = TensorBoard(log_dir="logs", histogram_freq=1)
model.fit(x_train, y_train, epochs=5, callbacks=[tf_callback])

model.evaluate(x_test,  y_test, verbose=2)

ログの格納先を指定して TensorBoard を表示します。

%tensorboard --logdir logs

実行結果

以下のようにノートブック中に TesorBoard を表示することができます。

現状では表示に時間はかかるようですが、 TensorBoard をノートブック上に表示することができました。