Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

TensorFlow 2.0 Keras: How to write image summaries for TensorBoard

I'm trying to setup an image recognition CNN with TensorFlow 2.0. To be able to analyze my image augmentation I'd like to see the images I feed into the network in tensorboard.

Unfortunately, I cannot figure out, how to do this with TensorFlow 2.0 and Keras. I also didn't really find documentation on this.

For simplicity, I'm showing the code of an MNIST example. How would I add the image summary here?

import tensorflow as tf
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()

def scale(image, label):
    return tf.cast(image, tf.float32) / 255.0, label

def augment(image, label):
    return image, label  # do nothing atm

dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset = dataset.map(scale).map(augment).batch(32)

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'])
model.fit(dataset, epochs=5, callbacks=[tf.keras.callbacks.TensorBoard(log_dir='D:\\tmp\\test')])
like image 747
andy Avatar asked Mar 29 '19 16:03

andy


People also ask

Can I use TensorBoard with keras?

Using TensorBoard with Keras Model.Start TensorBoard through the command line or within a notebook experience. The two interfaces are generally the same. In notebooks, use the %tensorboard line magic. On the command line, run the same command without "%".

How do you visualize a model in TensorBoard?

Select the Graphs dashboard by tapping “Graphs” at the top. You can also optionally use TensorBoard. dev to create a hosted, shareable experiment. By default, TensorBoard displays the op-level graph.

What is Logdir in TensorBoard?

--logdir is the directory you will create data to visualize. Files that TensorBoard saves data into are called event files. Type of data saved into the event files is called summary data. Optionally you can use --port=<port_you_like> to change the port TensorBoard runs on.


1 Answers

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)],
    )
like image 98
Szymon Maszke Avatar answered Sep 29 '22 15:09

Szymon Maszke