Use TensorBoard in Amazon SageMaker Studio

The following doc outlines how to install and run TensorBoard in Amazon SageMaker Studio.

Prerequisites

This tutorial requires an Amazon SageMaker Studio Domain. For more information, see Onboard to Amazon SageMaker Domain

Set Up TensorBoardCallback

1. Launch Studio, and open the Launcher. For more information, see Use the Amazon SageMaker Studio Launcher

2. In the Amazon SageMaker Studio Launcher, under Notebooks and compute resources, choose the Change environment button.

3. On the Change environment dialog, use the dropdown menus to select the TensorFlow 2.3 Python 3.7(optimized for CPU) Studio Image.

4. Back to the Launcher, click the Create notebook tile. Your notebook launches and opens in a new Studio tab.

5. Run this code from within your notebook cells.

6. Import the required packages.

   import os
   import datetime
   import tensorflow as tf

7. Create a Keras model.

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

8. Create a directory for your TensorBoard logs

   LOG_DIR = os.path.join(os.getcwd(), "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))

9. Run training with TensorBoard.

   model = create_model()
   model.compile(optimizer='adam',
                 loss='sparse_categorical_crossentropy',
                 metrics=['accuracy'])
                 
                 
   tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR, histogram_freq=1)
   
   model.fit(x=x_train,
             y=y_train,
             epochs=5,
             validation_data=(x_test, y_test),
             callbacks=[tensorboard_callback])

10. Generate the EFS path for the TensorBoard logs. You use this path to set up your logs from the terminal.

    EFS_PATH_LOG_DIR = "/".join(LOG_DIR.strip("/").split('/')[1:-1])
    print (EFS_PATH_LOG_DIR)

    Retrieve the EFS_PATH_LOG_DIR. You will need it in the TensorBoard installation section.

Install TensorBoard

1. Click on the  Amazon SageMaker Studio button on the top left corner of Studio to open the Amazon SageMaker Studio Launcher. This launcher must be opened from your root directory. For more information, see Use the Amazon SageMaker Studio Launcher

2. In the Launcher, under Utilities and files, click System terminal.

3. From the terminal, run the following commands. Copy EFS_PATH_LOG_DIR from the Jupyter notebook. You must run this from the /home/sagemaker-user root directory.

   pip install tensorboard
   tensorboard --logdir <EFS_PATH_LOG_DIR>

Launch TensorBoard

1. To launch TensorBoard, copy your Studio URL and replace lab? with proxy/6006/ as follows. You must include the trailing / character.

   https://<YOUR_URL>.studio.region.sagemaker.aws/jupyter/default/proxy/6006/

2. Navigate to the URL to examine your results.