[PYTHON] Easy 3 minutes TensorBoard in Google Colab (using TensorFlow 2.x)

When using TensorFlow, TensorBoard, which visualizes the state of learning, is often used. It is very convenient to be able to see the status interactively on the web screen, but it took some tricks to refer to it from within Google Colaboratory. TensorFlow 2.x seems to be easy to use on Google Colab. So I'll give it a try.

Execution environment

Use Google Colaboratory.

Sample code

You can read more about using TensorBoard from Google Colab on the official TensorFlow page. https://www.tensorflow.org/tensorboard/tensorboard_in_notebooks

Run the magic command of using 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

Then run the magic command to load the TensorBoard.

# Load the TensorBoard notebook extension
%load_ext tensorboard

Create a simple model using MNIST (image data of numbers often used as a sample). keras is very simple.

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'])

Give model.fit () a callback function for TensorBoard. This area is the same as how to use TensorBoard normally.

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)

Display TensorBoard by specifying the log storage location.

%tensorboard --logdir logs

Execution result

You can display TesorBoard in your notebook as follows:

It seems that it takes a long time to display at present, but I was able to display TensorBoard on my notebook.