Training a model with ArcFace layer according to code by 4uiiurz1 compatible with TensorFlow Keras

Question

I am trying to train a model with the ArcFace code taken from https://github.com/4uiiurz1/keras-arcface in which I took the ArcFace layer and added it to the model.

I created a small dataset of 4 samples, each is a (T_i, K) tensor with float32 and a label 0 or 1. Here K is constant and T_i changes for each sample.

The model is built as follows:

import tensorflow as tf
# from keras import layers as tfl
import tensorflow.keras.layers as tfl
from keras import Model

NUMBER_OF_CLASSES = 2
ARCFACE_SCALE = 30.0
ARCFACE_MARGIN = 0.5

input_tensor = tfl.Input(shape=[None, K], name="input")
label = tfl.Input(shape=(1,), name="labels")

output_tensor = # some layers applied to input tensor

arcface_layer = ArcFace(NUMBER_OF_CLASSES, s=ARCFACE_SCALE, m=ARCFACE_MARGIN) 
output_tensor = arcface_layer([output_tensor, label])
model = Model([input_tensor, label], output_tensor)

The model compilation:

model.compile(optimizer= 'adam' #optimizer_awd,
              loss='categorical_crossentropy',
              metrics='accuracy',
              loss_weights=None,
              weighted_metrics=None,
              run_eagerly=None,
              steps_per_execution=None,
              jit_compile=None,
              )

Up to this stage, it works successfully (doesn't throw an error).

The dataset is train_ds a TensorFlow.DataSet with features (T_i, K) and labels.

Now, after compiling the model I run the fit, according to the example in 4uiiurz1/keras-arcface readme, I need to enter

    model.fit(
        x=[x_train, y_train], 
        y=y_train,
        batch_size=1,
        epochs=3,
        verbose=1, #'auto',
        callbacks=None,
        validation_split=0.0,
        validation_data=([x_train, y_train], y_train),
        shuffle=True,
        class_weight=None,
        sample_weight=None,
        initial_epoch=0,
        steps_per_epoch=1, # steps_per_epoch,
        validation_steps=None,
        validation_batch_size=None,
        validation_freq=1,
        max_queue_size=None, # 10,
        workers=1,
        use_multiprocessing=False
    )

However, whatever object type I insert as x_train, y_train I get an error.

I tried:

    # prepare dataset
    x_train = []
    y_train = []
    for i in range(len(list(train_ds))):
        for j in range(1):
            x_train.append(list(train_ds)[i][0]) #.numpy())
            y_train.append(list(train_ds)[i][1]) #.numpy().astype(np.int16))


    # x_train = tf. convert_to_tensor(np.array(x_train))
    # y_train = tf. convert_to_tensor(np.array(y_train))

where train_ds is a TensorFlow.Dataset. Now it throws the following error:

Exception has occurred: IndexError
list index out of range

I also tried:

x_train, y_train = zip(*train_ds)

and got the same error.

I tried to increase the number of training examples from 4 to 102, and still - the same error repeats. When debugging the model.fit on dataset of 4 examples, we saw it also tries to access a fifth element, which doesn't exist and then it throws the error. Why does it try to access example of n+1 in the dataset of n examples?

So how do I convert a TensorFlow.DataSet to an object or list which will be compatible with model.fit and the 4uiiurz1/keras-arcface example?

Versions:

Python 3.9.5
TensforFlow 2.8.0
Keras 2.8.0

Cross-posted on StackOverflow: https://stackoverflow.com/questions/72868610/training-a-model-with-arcface-layer-according-to-code-by-4uiiurz1-compatible-wit