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 generatortf.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)],
)
