Speeding up reading TFRecords and Feeding them into Keras model

I would like to feed TFRecords into my model at a super fast rate. However, currently, my GPU(Single K80 on GCP) is at 0% load which is super slow on CloudML.

I have TFRecords in GCS: train_directory = gs://bucket/train/*.tfrecord, (around 100 files of 30mb-800mb in size), but for some reason it struggles to feed the data into my model fast enough for GPU.

Interestingly, loading data into memory and using numpy arrays using fit_generator() is 7x faster. There I can specify multi-processing and multi workers.

My current set up parses tf records and loads an infinite tf.Dataset. Ideally, the solution would save/prefecth some batches in memory, for the gpu to use on demand.

def _parse_func(record):
    """ Parses TF Record"""
    return x, y, w

def input_fn(f=filenames : list, b=BATCH_SIZE):
    """ Create tf.Dataset """
    d = tf.data.TFRecordDataset(filenames=f)
    d = d.map(map_func=_parse_func)
    d = d.repeat()
    d = d.batch(batch_size=b)
    d = d.prefetch(b * 100) # Prefetch 100 batches
    return d

Get train data

# get file names from GCS bucket and load them
train_files = file_io.list_directory(train_directory)
train_data = input_fn(f=train_files, b=BATCH_SIZE)

Fit the model

model.fit(x=train_data.make_one_shot_iterator())

I am running it on CloudML so GCS and CloudML should be pretty fast.

from Speeding up reading TFRecords and Feeding them into Keras model