[FIXED] Fit tf.data.Dataset with from_generator

Issue

I am training a NN to read 4 images and predict the next 4. Since I have a huge dataset, I have written this generator:

def my_gen(array_passed):
    for i in range(len(array_passed)-8):
        x = arr_train_ids[i:i+4]
        x = obtain_data(x) #shape (4,480,480,1)

        y = arr_train_ids[i+4:i+8]
        y = obtain_data(y) #same shape as x
        yield x, y

Where obtain_data is just open the file, load 4 480x480 matrices and reshape to (4, 480, 480, 1). Then, following a post from here, I do:

def my_input_fn(array_passed):
    dataset = tf.data.Dataset.from_generator(lambda: my_gen(array_passed),
                   output_types=(tf.float32, tf.float32),
                   ) #using TF2.3, output_signature not valid
    dataset = dataset.batch(1)
    #dataset = dataset.repeat(10)
    return dataset

train_dataset = my_input_fn(arr_train_ids)
val_dataset = my_input_fn(arr_val_ids)
test_dataset = my_input_fn(arr_test_ids)

I am not sure about the necessity of using repeat... It is now commented in my code.

The model, called rnc, goes here (omitted for brevity, ConvLSTM and Conv3D involved).

model = rnc(input_shape=(4, 480, 480, 1))
model.compile(optimizer=tf.keras.optimizers.Adam(lr=3e-4),loss='mae')
history = model.fit(train_dataset, validation_data=val_dataset, epochs=100)

I want to train 1 batch of data. In theory, 1 batch of data should be 4 matrices 480x480 with 1 channel, (4, 480, 480, 1) as input and also (4, 480, 480, 1) as output.

But I don't obtain anythinig from the training and I am quite suspicious. I wonder if the preparation of the TF dataset is something I am doing correctly.

Solution

If you checked the resulting input_batch and output_batch from your train_dataset, it has got to work.
In my simple reproductive example,

import tensorflow as tf
import tensorflow.keras as keras

x = tf.random.normal((100, 4, 480, 480, 1)) # Say I have 100 data.
y = tf.random.normal((100, 4, 480, 480, 1))

train_ds = tf.data.Dataset.from_tensor_slices((x, y))
train_ds = train_ds.batch(batch_size=32, drop_remainder=True)

# check whether this dataset really produces the things I want
sample_input_batch, sample_output_batch = next(iter(train_ds))
print(sample_input_batch.shape) # (32, 4, 480, 480, 1)
print(sample_output_batch.shape) # (32, 4, 480, 480, 1)

simple_model = keras.Sequential([
    keras.layers.InputLayer(input_shape=(4, 480, 480, 1)),
    keras.layers.Dense(10, activation='relu'),
    keras.layers.Dense(1, activation='relu')
])

simple_model.compile(loss=keras.losses.MeanSquaredError(),
                    optimizer=keras.optimizers.Adam(),
                    metrics=['mse'])

# check whether the model really produces the thing I expect
sample_predicted_batch = simple_model(sample_input_batch)
print(sample_predicted_batch.shape) # (32, 4, 480, 480, 1)

simple_model.fit(train_ds)
# 3/3 [==============================] - 0s 24ms/step - loss: 1.1743 - mse: 1.1743
# Then it should work!

Besides, if you have a huge dataset, you don't really need to use repeat() method of tf.data.API.

Moreover, you've got to specify steps_per_epoch argument to some number if you are using repeat() method otherwise it means to run forever. You can check about steps_per_epoch in here.

Answered By - HyeonPhil Youn