TensorFlow 2.0 Keras: How to write image summaries for TensorBoard

Except providing an answer to your question I will make the code more TF2.0-like. If you have any questions/need clarification, please post a comment down below.

1. Loading data

I would advise to use Tensorflow Datasets library. There is absolutely no need to load data in numpy and transform it to tf.data.Dataset if one can do it in a single line:

import tensorflow_datasets as tfds

dataset = tfds.load("mnist", as_supervised=True, split=tfds.Split.TRAIN)

Line above will only return TRAIN split (read more about those here).

2. Define Augmentations and Summaries

In order to save images, one has to keep tf.summary.SummaryWriter object throughout each pass.

I have created a convenient wrapping class with __call__ method for easy usage with tf.data.Dataset's map capabilities:

import tensorflow as tf

class ExampleAugmentation:
    def __init__(self, logdir: str, max_images: int, name: str):
        self.file_writer = tf.summary.create_file_writer(logdir)
        self.max_images: int = max_images
        self.name: str = name
        self._counter: int = 0

    def __call__(self, image, label):
        augmented_image = tf.image.random_flip_left_right(
            tf.image.random_flip_up_down(image)
        )
        with self.file_writer.as_default():
            tf.summary.image(
                self.name,
                augmented_image,
                step=self._counter,
                max_outputs=self.max_images,
            )

        self._counter += 1
        return augmented_image, label

name will be the name under which each part of images will be saved. Which part you may ask - the part defined by max_outputs.

Say image in __call__ will have shape (32, 28, 28, 1), where the first dimension is batch, second width, third height and last channels (in case of MNIST only onel but this dimension is needed in tf.image augmentations). Furthermore, let's say max_outputs is specified as 4. In this case, only 4 first images from batch will be saved. Default value is 3, so you may set it as BATCH_SIZE to save every image.

In Tensorboard, each image will be a separate sample over which you can iterate at the end.

_counter is needed so the images will not be overwritten (I think, not really sure, clarification from someone else would be nice).

Important: You may want to rename this class to something like ImageSaver when doing more serious buisness and move augmentation to separate functors/lambda functions. It suffices for presentation purposes I guess.

3. Setup global variables

Please do not mix function declaration, global variables, data loading and others (like loading data and creating function afterwards). I know TF1.0 encouraged this type of programming but they are trying to get away from it and you might want to follow the trend.

Below I have defined some global variables which will be used throughout next parts, pretty self-explanatory I guess:

BATCH_SIZE = 32
DATASET_SIZE = 60000
EPOCHS = 5

LOG_DIR = "/logs/images"
AUGMENTATION = ExampleAugmentation(LOG_DIR, max_images=4, name="Images")

4. Dataset augmentation

Similar to yours but with a little twist:

dataset = (
    dataset.map(
        lambda image, label: (
            tf.image.convert_image_dtype(image, dtype=tf.float32),
            label,
        )
    )
    .batch(BATCH_SIZE)
    .map(AUGMENTATION)
    .repeat(EPOCHS)
)

• repeat is needed as the loaded dataset is a generator
• tf.image.convert_image_dtype - better and more readable option than explicit tf.cast mixed with division by 255 (and ensures proper image format)
• batching done before augmentation just for the sake of presentation

5. Define model, compile, train

Almost as you did in your example, but I have provided additional steps_per_epoch, so fit knows how many batches constitute an epoch:

model = tf.keras.models.Sequential(
    [
        tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
        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"]
)
model.fit(
    dataset,
    epochs=EPOCHS,
    steps_per_epoch=DATASET_SIZE // BATCH_SIZE,
    callbacks=[tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR)],
)

Not much to explain other than that I think.

6. Run Tensorboard

Since TF2.0 one can do it inside colab using %tensorboard --logdir /logs/images, just wanted to add this for others who may visit this issue. Do it however you like, anyways you know how to do it for sure.

Images should be inside IMAGES and each sample named by name provided to AUGMENTATION object.

7. Whole code (to make everyone's life easier)

import tensorflow as tf
import tensorflow_datasets as tfds


class ExampleAugmentation:
    def __init__(self, logdir: str, max_images: int, name: str):
        self.file_writer = tf.summary.create_file_writer(logdir)
        self.max_images: int = max_images
        self.name: str = name
        self._counter: int = 0

    def __call__(self, image, label):
        augmented_image = tf.image.random_flip_left_right(
            tf.image.random_flip_up_down(image)
        )
        with self.file_writer.as_default():
            tf.summary.image(
                self.name,
                augmented_image,
                step=self._counter,
                max_outputs=self.max_images,
            )

        self._counter += 1
        return augmented_image, label


if __name__ == "__main__":

    # Global settings

    BATCH_SIZE = 32
    DATASET_SIZE = 60000
    EPOCHS = 5

    LOG_DIR = "/logs/images"
    AUGMENTATION = ExampleAugmentation(LOG_DIR, max_images=4, name="Images")

    # Dataset

    dataset = tfds.load("mnist", as_supervised=True, split=tfds.Split.TRAIN)

    dataset = (
        dataset.map(
            lambda image, label: (
                tf.image.convert_image_dtype(image, dtype=tf.float32),
                label,
            )
        )
        .batch(BATCH_SIZE)
        .map(AUGMENTATION)
        .repeat(EPOCHS)
    )

    # Model and training

    model = tf.keras.models.Sequential(
        [
            tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
            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"]
    )
    model.fit(
        dataset,
        epochs=EPOCHS,
        steps_per_epoch=DATASET_SIZE // BATCH_SIZE,
        callbacks=[tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR)],
    )