Reading .h5 file is extremely slow

Question

My data is stored in .h5 format. I use a data generator to fit the model and it is extremely slow. A snippet of my code is provided below.

def open_data_file(filename, readwrite="r"):
    return tables.open_file(filename, readwrite)

data_file_opened = open_data_file(os.path.abspath("../data/data.h5"))

train_generator, validation_generator, n_train_steps, n_validation_steps = get_training_and_validation_generators(
        data_file_opened,
        ......)

where:

def get_training_and_validation_generators(data_file, batch_size, ...):
    training_generator = data_generator(data_file, training_list,....)

data_generator function is as follows:

def data_generator(data_file, index_list,....):
      orig_index_list = index_list
    while True:
        x_list = list()
        y_list = list()
        if patch_shape:
            index_list = create_patch_index_list(orig_index_list, data_file, patch_shape,
                                                 patch_overlap, patch_start_offset,pred_specific=pred_specific)
        else:
            index_list = copy.copy(orig_index_list)

        while len(index_list) > 0:
            index = index_list.pop()
            add_data(x_list, y_list, data_file, index, augment=augment, augment_flip=augment_flip,
                     augment_distortion_factor=augment_distortion_factor, patch_shape=patch_shape,
                     skip_blank=skip_blank, permute=permute)
            if len(x_list) == batch_size or (len(index_list) == 0 and len(x_list) > 0):
                yield convert_data(x_list, y_list, n_labels=n_labels, labels=labels, num_model=num_model,overlap_label=overlap_label)
                x_list = list()
                y_list = list()

add_data() is as follows:

def add_data(x_list, y_list, data_file, index, augment=False, augment_flip=False, augment_distortion_factor=0.25,
             patch_shape=False, skip_blank=True, permute=False):
    '''
    add qualified x,y to the generator list
    '''
#     pdb.set_trace()
    data, truth = get_data_from_file(data_file, index, patch_shape=patch_shape)
    
    if np.sum(truth) == 0:
        return
    if augment:
        affine = np.load('affine.npy')
        data, truth = augment_data(data, truth, affine, flip=augment_flip, scale_deviation=augment_distortion_factor)

    if permute:
        if data.shape[-3] != data.shape[-2] or data.shape[-2] != data.shape[-1]:
            raise ValueError("To utilize permutations, data array must be in 3D cube shape with all dimensions having "
                             "the same length.")
        data, truth = random_permutation_x_y(data, truth[np.newaxis])
    else:
        truth = truth[np.newaxis]

    if not skip_blank or np.any(truth != 0):
        x_list.append(data)
        y_list.append(truth)

Model training:

def train_model(model, model_file,....):
    model.fit(training_generator,
                        steps_per_epoch=steps_per_epoch,
                        epochs=n_epochs,
                        verbose = 2,
                        validation_data=validation_generator,
                        validation_steps=validation_steps)

My dataset is large: data.h5 is 55GB. It takes around 7000s to complete one epoch. And I get a segmentation fault error after like 6 epochs. The batch size is set to 1, because otherwise, I get a resource exhausted error. Is there an efficient way to read data.h5 in the generator so that training is faster and doesn't lead to out-of-memory errors?

Answer 1

This is the start of my answer. I looked at your code, and you have a lot of calls to read the .h5 data. By my count, the generator makes 6 read calls for every loop on training_list and validation_list. So, that's almost 20k calls on ONE training loop. It's not clear (to me) if the generators are called on every training loop. If they are, multiply by 2268 loops.

Efficiency of HDF5 file read depends on the number to calls to read the data (not just the amount of data). In other words, it is faster to read 1GB of data in a single call than it is to read the same data with 1000 calls x 1MB at a time. So the first thing we need to determine is the amount of time spent reading data from the HDF5 file (to be compare to your 7000s).

I isolated the PyTables calls that read the data file. From that, I built a simple program that mimics the behavior of your generator function. Currently it makes a single training loop on the entire sample list. Increase n_train and n_epoch values if you want the to run a longer test. (Note: The code syntax is correct. However without the file, so can't verify the logic. I think it's correct, but you may have to fix small errors.)

See code below. It should run standalone (all dependencies are imported). It prints basic timing data. Run it to benchmark your generator.

import tables as tb
import numpy as np
from random import shuffle 
import time

with tb.open_file('../data/data.h5', 'r') as data_file:

    n_train = 1
    n_epochs = 1
    loops = n_train*n_epochs
    
    for e_cnt in range(loops):  
        nb_samples = data_file.root.truth.shape[0]
        sample_list = list(range(nb_samples))
        shuffle(sample_list)
        split = 0.80
        n_training = int(len(sample_list) * split)
        training_list = sample_list[:n_training]
        validation_list = sample_list[n_training:]
        
        start = time.time()
        for index_list in [ training_list, validation_list ]:
            shuffle(index_list)
            x_list = list()
            y_list = list()
            
            while len(index_list) > 0:
                index = index_list.pop() 
                
                brain_width = data_file.root.brain_width[index]
                x = np.array([modality_img[index,0,
                                           brain_width[0,0]:brain_width[1,0]+1,
                                           brain_width[0,1]:brain_width[1,1]+1,
                                           brain_width[0,2]:brain_width[1,2]+1] 
                              for modality_img in [data_file.root.t1,
                                                   data_file.root.t1ce,
                                                   data_file.root.flair,
                                                   data_file.root.t2]])
                y = data_file.root.truth[index, 0,
                                         brain_width[0,0]:brain_width[1,0]+1,
                                         brain_width[0,1]:brain_width[1,1]+1,
                                         brain_width[0,2]:brain_width[1,2]+1]    
                
                x_list.append(data)
                y_list.append(truth)
    
        print(f'For loop:{e_cnt}')
        print(f'Time to read all data={time.time()-start:.2f}')