Largo tiempo de inicialización para model.fit cuando se usa el conjunto de datos de tensorflow del generador

Esta es mi primera pregunta sobre el desbordamiento de pila. Me disculpo de antemano por el mal formato y la sangría debido a mis problemas con la interfaz.

Especificaciones del entorno:

Versión de Tensorflow: GPU 2.7.0 (probada y funcionando correctamente)

Versión de Python - 3.9.6

CPU: Intel Core i7 7700HQ

GPU: NVIDIA GTX 1060 de 3 GB

RAM - 16GB DDR4 2400MHz

HDD - 1 TB 5400 RPM

Enunciado del problema:

Deseo entrenar un modelo TensorFlow 2.7.0 para realizar una clasificación multietiqueta con seis clases en tomografías computarizadas almacenadas como imágenes DICOM. El conjunto de datos es de Kaggle, vínculo aquí. Las etiquetas de entrenamiento se almacenan en un archivo CSV y los nombres de las imágenes DICOM tienen el formato ID_"caracteres aleatorios".dcm. Las imágenes tienen un tamaño combinado de 368 GB.

Enfoque utilizado:

1. El archivo CSV que contiene las etiquetas se importa a pandas DataFrame y los nombres de archivo de imagen se establecen como índice.

2. Se crea un generador de datos simple para leer la imagen DICOM y la etiquetas iterando en las filas del DataFrame. Este generador se utiliza para crear un conjunto de datos de entrenamiento usando tf.data.Dataset.from_generator. las imagenes son preprocesado usando bsb_window().

3. El conjunto de datos de entrenamiento se baraja y se divide en un entrenamiento (90%) y conjunto de validación (10%)

4. El modelo se crea con Keras Sequential, se compila y ajusta con los conjuntos de datos de entrenamiento y validación creados anteriormente.

código:

def train_generator():
   for row in df.itertuples():
       image = pydicom.dcmread(train_images_dir + row.Index + ".dcm")
       try:
           image = bsb_window(image)
       except:
           image = np.zeros((256,256,3))
       labels = row[1:]
       yield image, labels

train_images = tf.data.Dataset.from_generator(train_generator, 
                                              output_signature = 
                                             ( 
                                               tf.TensorSpec(shape = (256,256,3)), 
                                               tf.TensorSpec(shape = (6,))
                                              )
                                              )
train_images = train_images.batch(4)
TRAIN_NUM_FILES = 752803
train_images = train_images.shuffle(40)
val_size = int(TRAIN_NUM_FILES * 0.1)
val_images = train_images.take(val_size)
train_images = train_images.skip(val_size)

def create_model():
   model = Sequential([
                       InceptionV3(include_top = False, input_shape = (256,256,3), weights = "imagenet"),
                       GlobalAveragePooling2D(name = "avg_pool"),
                       Dense(6, activation = "sigmoid", name = "dense_output"),
                       ])
   model.compile(loss = "binary_crossentropy", 
                 optimizer = tf.keras.optimizers.Adam(5e-4), 
                 metrics = ["accuracy", tf.keras.metrics.SpecificityAtSensitivity(0.8)]
                 )
   return model

model = create_model()
history = model.fit(train_images, 
                    batch_size=4, 
                    epochs=5, 
                    verbose=1, 
                    validation_data=val_images
                    )

Problema:

Al ejecutar este código, hay un retraso de algunas horas de alto uso del disco (~30 MB/s de lecturas) antes de que comience el entrenamiento. Cuando se crea un DataGenerator usando tf.keras.utils.Sequence, el entrenamiento comienza segundos después de llamar a model.fit().

Posibles causas:

1. Iterando sobre un DataFrame de pandas en train_generator(). No estoy seguro de cómo evitar este problema.
2. El uso de funciones externas para preprocesar y cargar los datos.
3. El uso de los métodos take() y skip() para crear conjuntos de datos de entrenamiento y validación.

¿Cómo optimizo este código para que se ejecute más rápido? Escuché que dividir el generador de datos en creación de etiquetas, funciones de preprocesamiento de imágenes y operaciones paralelas mejoraría el rendimiento. Aún así, no estoy seguro de cómo aplicar esos conceptos en mi caso. Cualquier consejo sería muy apreciado.