Training a TensorFlow-Keras model with extra layer which gets also the labels as input

Question

I want to train a model which create a features vector for a RGB image (2D array with 3 channels), and using that features vector, a classifier will decide what to do (e.g. person recognition from an image, assign a label by choosing the "closest" pre-trained features vectors (of the people enrolled to the system). To do so I use a categorical cross-entropy. In the training phase I apply categorical softmax on the features vector as an extra layer, and get as output the probability to be in each label or class, then I use the softmax output and the training label to compute the loss. So, for working or testing, the model receives just one input: the image, and outputs a features vector. While for training the model receives pairs: the image and its label.

I want to train such a model, with pair [image,label] input in the training phase, and [image] input in the testing or working phase.

I use TensorFlow 2.8 and Keras 2.8 with Python 3.9.5.

The code (with a toy model and some random data):

# ==============================================================================
# Imports
# ==============================================================================
import numpy as np

import tensorflow as tf
import keras
import keras.backend as K
from keras import layers as tfl
from keras import Model

# ==============================================================================
# Switch case layer, behaves differently for training and testing
# ==============================================================================
class Switch(tf.keras.layers.Layer):
    def __init__(self, **kwargs):

        super().__init__(**kwargs)

    def call(self, inputs, training=None):
        x = tf.identity(inputs)
        if training:
            y = tfl.Input(shape=(2,), name="label")
            output_tensor = tf.nn.softmax_cross_entropy_with_logits(y, x)
            return output_tensor
        else:
            output_tensor = tf.identity(x, name="output")
            return output_tensor

# ==============================================================================
# Define model
# ==============================================================================
inputs = keras.Input(shape=(4, 4, 3))
conv = keras.layers.Conv2D(filters=2, kernel_size=2)(inputs)
pooling = keras.layers.GlobalAveragePooling2D()(conv)
feature = keras.layers.Dense(10)(pooling)
outputs = Switch()(feature)
# output = tf.identity(feature)

model = keras.Model(inputs, outputs)

# ==============================================================================
# Training data
# ==============================================================================
tf.random.set_seed(42)
x_train = tf.random.normal((5, 4, 4, 3))
y_train = tf.constant([1, 1, 0, 0, 2])

# ==============================================================================
# Train model
# ==============================================================================
model.compile(optimizer='adam',
                  loss='categorical_crossentropy',
                  metrics='accuracy',)
model.fit(
    x=x_train, 
    y=y_train,
    epochs=3,
    verbose='auto',
    shuffle=True,
    initial_epoch=0,
    max_queue_size=10
)

The Switch layer is based on: Is it possible to add different behavior for training and testing in keras Functional API

If I understand correctly, when using model.fit, the model's call is automatically invoked with training=True.

However, when I run the model I get the following error:

TypeError: You are passing KerasTensor(type_spec=TensorSpec(shape=(), dtype=tf.float32, name=None), name='Placeholder:0', description="created by layer 'tf.cast_4'"), an intermediate Keras symbolic input/output, to a TF API that does not allow registering custom dispatchers, such as tf.cond, tf.function, gradient tapes, or tf.map_fn. Keras Functional model construction only supports TF API calls that do support dispatching, such as tf.math.add or tf.reshape. Other APIs cannot be called directly on symbolic Kerasinputs/outputs. You can work around this limitation by putting the operation in a custom Keras layer call and calling that layer on this symbolic input/output.

When I pass:

model.fit(
    x=[x_train, y_train], 
    y=y_train,

I receive the following error:

ValueError: Layer "model" expects 1 input(s), but it received 2 input tensors. Inputs received: [, ]

The problem is probably due to the Switch layer.

How do I solve it and how I train a model in which the training phase input and output are different than in the working phase (gets image, outputs features vector)?

Training a TensorFlow-Keras model with extra layer which gets also the labels as input

0 Answers