Входной конвейер набора данных Tensorflow-IO с очень большими файлами HDF5

Question

У меня очень большие обучающие (30 ГБ) файлы.
Поскольку все данные не помещаются в моей доступной ОЗУ, я хочу прочитать данные по пакетам.
Я увидел, что есть пакет Tensorflow-io, который реализовал способ для чтения HDF5 в Tensorflow таким образом благодаря функции tfio.IODataset.from_hdf5()
Затем, поскольку tf.keras.model.fit() принимает tf.data.Dataset в качестве входных данных, содержащих как сэмплы, так и цели, мне нужно сжать мой X и Y вместе, а затем используйте .batch and .prefetch для загрузки в память только необходимых данных. Для тестирования я пытался применить этот метод к меньшим выборкам: обучение (9 ГБ), проверка (2,5 ГБ) и тестирование (1,2 ГБ), которые, как я знаю, хорошо работают, потому что они могут вписаться в память, и у меня хорошие результаты (точность 70% и 1 потеря).
Учебные файлы хранятся в файлах HDF5, разбитых на файлы сэмплов (X) и меток (Y), например:

X_learn.hdf5  
X_val.hdf5  
X_test.hdf5  
Y_test.hdf5  
Y_learn.hdf5  
Y_val.hdf5

Вот мой код:

BATCH_SIZE = 2048
EPOCHS = 100

# Create an IODataset from a hdf5 file's dataset object  
x_val = tfio.IODataset.from_hdf5(path_hdf5_x_val, dataset='/X_val')
y_val = tfio.IODataset.from_hdf5(path_hdf5_y_val, dataset='/Y_val')
x_test = tfio.IODataset.from_hdf5(path_hdf5_x_test, dataset='/X_test')
y_test = tfio.IODataset.from_hdf5(path_hdf5_y_test, dataset='/Y_test')
x_train = tfio.IODataset.from_hdf5(path_hdf5_x_train, dataset='/X_learn')
y_train = tfio.IODataset.from_hdf5(path_hdf5_y_train, dataset='/Y_learn')

# Zip together samples and corresponding labels
train = tf.data.Dataset.zip((x_train,y_train)).batch(BATCH_SIZE, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)
test = tf.data.Dataset.zip((x_train,y_train)).batch(BATCH_SIZE, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)
val = tf.data.Dataset.zip((x_train,y_train)).batch(BATCH_SIZE, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)

# Build the model
model = build_model()

# Compile the model with custom learing rate function for Adam optimizer
model.compile(loss='categorical_crossentropy',
               optimizer=Adam(lr=lr_schedule(0)),
               metrics=['accuracy'])

# Fit model with class_weights calculated before
model.fit(train,
          epochs=EPOCHS,
          class_weight=class_weights_train,
          validation_data=val,
          shuffle=True,
          callbacks=callbacks)

Этот код выполняется, но потери очень высоки (300+) и точность падает с нуля до 0 (0,30 -> 4 * e ^ -5) с самого начала ... Я не понимаю, что я делаю неправильно, я что-то упустил?