VGG16 preprocessing dataset generator to dataset mapping

I have a VGG16 model implemented with Keras/tensorflow.

When I call model.fit, I pass in a generator of data. The generator does transforms necessary for a VGGNet:

1. Preprocess the images with vgg16.preprocess_input
2. Convert the label to a one-hot vector via to_categorical

The generator can be seen below and works. Unfortunately, since there are multiple epochs, I have to set dataset.repeat(-1) (infinitely repeat) so the generator doesn't run out. This in turn requires one to pass a steps_per_epoch so a given iteration of training can complete. As you're probably thinking, this is brittle, (hinges on a known dataset cardinality)!

I have decided it's best to preprocess the training Dataset once up front using Dataset.map. However, I am struggling with the construction of a mapping function, it seems to_categorical doesn't work with a tf.Tensor. Down below is what I have right now, but I am not sure if there's a latent bug.

How can I correctly translate the below Dataset generator into a Dataset.map function?

Current Dataset Generator

This is implemented (and known to work) with Python 3.8 and tensorflow==2.4.4.

from typing import Iterable, Tuple

import numpy as np
import tensorflow as tf


def make_vgg_preprocessing_generator(
    dataset: tf.data.Dataset, num_repeat: int = -1
) -> Iterable[Tuple[tf.Tensor, np.ndarray]]:
    num_classes = len(dataset.class_names)
    for batch_images, batch_labels in dataset.repeat(num_repeat):
        pre_images = tf.keras.applications.vgg16.preprocess_input(batch_images)
        pre_labels = tf.keras.utils.to_categorical(batch_labels, num_classes)
        yield pre_images, pre_labels


train_ds: tf.data.Dataset  # Not provided in this sample
model.fit(
    make_vgg_preprocessing_generator(train_ds)
    epochs=10,
    steps_per_epoch=10,  # Required since default num_repeat is indefinitely
)

Dataset.map Function

Here is my current translation that I would like to improve.

def vgg_preprocess_dataset(dataset: tf.data.Dataset) -> tf.data.Dataset:
    num_classes = len(dataset.class_names)

    def _preprocess(x: tf.Tensor, y: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        pre_x = tf.keras.applications.vgg16.preprocess_input(x)
        pre_y = tf.one_hot(y, depth=num_classes)
        return pre_x, pre_y

    return dataset.map(_preprocess)

Yes, you're on the right track! You'll want to replace to_categorical with tf.one_hot, just as you have, as tf.one_hot is specifically for tensors, and is designed for this context. Next, you might want to play around with some of the other tf.data.Dataset methods here and add them to your pipeline. Right now, your batch size will be one sample, and un-shuffled. An example of some other processing you might do:

def vgg_preprocess_dataset(dataset: tf.data.Dataset, batch_size=32, shuffle_buffer=1000) -> tf.data.Dataset:
    num_classes = len(dataset.class_names)

    def _preprocess(x: tf.Tensor, y: tf.Tensor):
        pre_x = tf.keras.applications.vgg16.preprocess_input(x)
        pre_y = tf.one_hot(y, depth=num_classes)
        # pre_y = to_categorical(y, num_classes)
        return pre_x, pre_y

    # bigger buffer is better but slower
    dataset = dataset.shuffle(shuffle_buffer) 

    # do your mapping after shuffle
    dataset = dataset.map(_preprocess)

    # next batch it
    dataset = dataset.batch(batch_size)
    
    # this allows your CPU to fetch the next batch (do the above shuffling, mapping, etc) during the 
    # current GPU pass, so that the GPU has minimal downtime
    dataset = dataset.prefetch(2)

    return dataset

ds = vgg_preprocess_dataset(ds)

# and you just pass it right to fit!
model.fit(ds)