Data augmentation for multiple output heads in Keras

Question

I have a transfer learning based two output classification problem. So, accordingly, I have formatted my data to have X_train as a (number of samples, height, width, channels) numpy array, y_train1 as (number of samples,) numpy array and y_train2 as (number of samples,) numpy array.

As I am not training using directory structure, I am using ImageDataGenerator.flow(). However, I am not able to figure out how I can pass two label arrays because, it is taking the labels as (2, number of samples) when I send it as [y_train1, y_train2] list.

I am able to train the network without Keras data augmentation (for two outputs). But, I am not able to apply data augmentation.

I am trying to do the following:

datagen = ImageDataGenerator(horizontal_flip=True, vertical_flip=True, rescale=1./255, class_mode="multi-label")

model.fit(datagen.flow(X_train, [y_train1, y_train2], batch_size=batch_size), batch_size=batch_size, epochs=nb_epochs, steps_per_epoch=spe, validation_data=(X_val, [y_val1, y_val2]))

Also, ImageDataGenerator.flow does not have class_mode unlike ImageDataGenerator.flow_from_dataframe.

Any suggestions/help would be appreciated!

References:

1. Data Augmentation Multi Outputs (No answer. I upvoted this just now)
2. Get multiple output from Keras (Does not explain data augmentation)

Answer 1

Please refer to the source code provided at https://gist.github.com/swghosh/f728fbba5a26af93a5f58a6db979e33e which should assist you in writing custom generators (basis ImageDataGenerator) for training end to end multi-output models. In the provided example, GoogLeNet is being trained which consists of two auxiliary classifiers and thus, comprising 3 outputs in the complete model.

The output from ImageDataGenerator.flow_from_directory() have been passed into a Python function with yield statements such that the targets can be repeated thrice as per requirement.

def three_way(gen):
    for x, y in gen:
        yield x, [y, y, y]

Also, please make sure to pass a value to steps_per_epoch argument in case of model.fit() or model.fit_generator() so as to ensure that the training loop doesn't run indefinitely. (This is because the Python generator object won't be inherited from Sequence class and number of batches/steps cannot be inferred implicitly by Keras)

Also, you may consider using an image data pipeline constructed using tf.data APIs combined with Random* Layers and Preprocessing Layers which is a more modern approach for tf.keras compared to the old ImageDataGenerator. There should be a bunch of performance and usability benefits in case you transition. Let me know if you are looking for an example.

Please look at this example in case you want to implement multi-output multi-label classification model. https://www.pyimagesearch.com/2018/06/04/keras-multiple-outputs-and-multiple-losses/

I'd still advise to use tf.data.Dataset instead of raw tensors which'd be more feasible for larger datasets that do not fit into memory plus, the tf.data APIs provides for an additional bunch of flexibilities.