Comparison of training time based on what you feed to the .fit method (np.array, ImageDataGenerator with different arguments, tf.data)

training_performance_inputs_type.py

"""
Performance comparison between ImageDataGenerator and tf.data.Dataset on a simple task

Dependencies available at https://www.github.com/sicara/tf2-evonorm
"""
import math
from pathlib import Path

import click
import tensorflow as tf
import tensorflow_addons as tfa

from scripts.resnet import ResnetBuilder, BATCH_NORM_NAME


@click.command()
@click.option("--dataset_name", default="cifar10", help="tf.keras dataset to use")
@click.option("--logdir", default="perf_logs", help="Logs directory")
def launch_training(dataset_name, logdir):
    dataset_module = getattr(tf.keras.datasets, dataset_name)
    (x_train, y_train), (x_test, y_test) = dataset_module.load_data()

    x_train = x_train.astype("float32") / 255.
    x_test = x_test.astype("float32") / 255.
    y_train = tf.keras.utils.to_categorical(y_train)
    y_test = tf.keras.utils.to_categorical(y_test)

    input_shape = x_train.shape[1:]
    num_classes = y_train.shape[1]
    batch_size = 64
    epochs = 15

    data_generator_without_rotate = tf.keras.preprocessing.image.ImageDataGenerator(
        horizontal_flip=True,
    )

    data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
        horizontal_flip=True,
        rotation_range=10,
    )

    def random_rotation(image, rotation_range):
        radian_rotation_range = rotation_range * math.pi / 180
        rotation = tf.random.uniform(shape=[], minval=-radian_rotation_range, maxval=radian_rotation_range)
        return tfa.image.rotate(image, rotation)

    dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).cache()
    dataset = dataset.map(lambda image, label: (tf.image.random_flip_left_right(image), label), num_parallel_calls=True)
    dataset = dataset.map(lambda image, label: (random_rotation(image, rotation_range=10), label), num_parallel_calls=True)
    dataset = dataset.shuffle(buffer_size=500).batch(batch_size).repeat()

    base_path = Path(f"{logdir}")
    np_array_path = base_path / "np.array"
    image_data_generator_path = base_path / "ImageDataGenerator"
    image_data_generator_without_rotate_path = base_path / "ImageDataGenerator.no_rotate"
    tf_data_path = base_path / "tf.data"

    model = ResnetBuilder.build_resnet_18(input_shape, num_classes, block_fn_name=BATCH_NORM_NAME)
    model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

    model.fit(
        x=x_train,
        y=y_train,
        batch_size=batch_size,
        steps_per_epoch=len(x_train) // batch_size,
        validation_data=(x_test, y_test),
        validation_steps=len(x_test) // batch_size,
        epochs=epochs,
        callbacks=[tf.keras.callbacks.TensorBoard(np_array_path)],
        shuffle=True,
    )  # 3'27" on RTX

    model = ResnetBuilder.build_resnet_18(input_shape, num_classes, block_fn_name=BATCH_NORM_NAME)
    model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

    model.fit(
        data_generator_without_rotate.flow(x_train, y_train, batch_size=batch_size),
        steps_per_epoch=len(x_train) // batch_size,
        validation_data=(x_test, y_test),
        validation_steps=len(x_test) // batch_size,
        epochs=epochs,
        callbacks=[tf.keras.callbacks.TensorBoard(image_data_generator_without_rotate_path)],
        shuffle=True,
    )  # 3'29" on RTX

    model = ResnetBuilder.build_resnet_18(input_shape, num_classes, block_fn_name=BATCH_NORM_NAME)
    model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

    model.fit(
        data_generator.flow(x_train, y_train, batch_size=batch_size),
        steps_per_epoch=len(x_train) // batch_size,
        validation_data=(x_test, y_test),
        validation_steps=len(x_test) // batch_size,
        epochs=epochs,
        callbacks=[tf.keras.callbacks.TensorBoard(image_data_generator_path)],
        shuffle=True,
    )  # 4'46" on RTX

    model = ResnetBuilder.build_resnet_18(input_shape, num_classes, block_fn_name=BATCH_NORM_NAME)
    model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

    model.fit(
        dataset,
        steps_per_epoch=len(x_train) // batch_size,
        validation_data=tf.data.Dataset.from_tensor_slices((x_test, y_test)).cache().batch(batch_size).repeat(),
        validation_steps=len(x_test) // batch_size,
        epochs=epochs,
        callbacks=[tf.keras.callbacks.TensorBoard(tf_data_path)],
        shuffle=True,
    )  # 3'27" on RTX


if __name__ == "__main__":
    launch_training()