Fast loading multiple .npy files into data generator

Question

I do have thousands of .npy files stored in my hard disk, each containing a single matrix with dimensions [128, T], where T is variable (on average T=800). Each .npy file has size around 2Mb, depending on the matrix shape.

These matrices are then passed to a generator, which yields batches of 32 to a neural network. The Python code used to pass the matrices into the generator is:

def load_batch(path_list):
    np_list = []
        
    for path in path_list:
        np_list.append(np.load(path))
           
    return np_list

which, given a list of paths of the .npy files, returns a list of the corresponding NumPy matrices.

This code takes, on average, 0.6s to return a list of 32 matrices. I am using append because this is usually a quick operation.

I am aware that the speed of the hard disk buffer does have an influence on timings but, right now, I really would like to shrink the amount of time required as much as possible by just modifying the code in a smart way.

As an alternative, I tried implementing multi-processing:

from multiprocessing import Pool
def reader(filename):
    return np.load(filename)

def load_multiprocess(path_list, n_cores=5):
    pool = Pool(n_cores)
    np_list = pool.map(reader, path_list)

    return np_list

However, the performance is much worse. I had a look around stackoverflow, and I got the idea that my specific application could not benefit from multiprocessing.

To summarize, I am looking for any kind of advice for one of these two tasks:

1. Improving the speed of the first code (even 0.1s less would mean a lot).
2. Using multiprocessing in the right way, if possible.

SOLUTION AND BENCHMARK

Out of the three methods here proposed, user7138814's solution seems to generally improve a lot the execution speed. However, things seem to change when the data is loaded while training a neural network: even though mapping is by itself still the quicker method for loading data, the overall training time seems to increase, I have no idea where and why, as timings using the mapping load are always better.

Below, I will do a benchmark of the three methods.First, define the methods:

import numpy as np

# my initial method
def load_batch(path_list):
    np_list = []
        
    for path in path_list:
        np_list.append(np.load(path))
           
    return np_list
    
# Aaj Kaal's method
def load_batch1(path_list):
    return [np.load(path) for path in path_list]
    
# user7138814's method    
def load_batch2(path_list):
    np_list = []
        
    for path in path_list:
        np_list.append(np.load(path, mmap_mode='r'))
    return np_list

I defined a list of paths as follows:

batches_list = []
batch_size = 32
for n in range(0,150):
    batches_list.append(X_path_list[n*batch_size:n*batch_size+batch_size])

The list contains 150 batches of 32 paths each, it should be enough to calculate the mean.

Then, each method is executed using passing to it exactly the same data.

import time

# my initial method
timing0 = []
for l in batches_list:
    start = time.time()
    load_batch(l)
    end = time.time()
    timing0.append(end-start)
print(np.mean(timing0))

# Aaj Kaal's method
timing1 = []
for l in batches_list:
    start = time.time()
    load_batch1(l)
    end = time.time()
    timing1.append(end-start)
print(np.mean(timing1))

# user7138814's method 
timing2 = []
for l in batches_list:
    start = time.time()
    load_batch2(l)
    end = time.time()
    timing2.append(end-start)
    print(np.mean(timing2))

Output (mean timing in seconds over 150 executions):

0.022530150413513184

0.022546884218851725

0.009580903053283692

Results seem to be consistent when changing length of batches_list and batch_size.

Answer 1

Maybe memory mapping the files will be beneficial due to lazy loading. If you would use for example

np.load(filename, mmap_mode='r')

the creation of the numpy array becomes almost a no-op, but later in the pipeline you pay the price. This could provide a speedup if it results in processing the data in parallel with reading from disk.

Answer 2

Did you try using use list comprehension. Replace

def load_batch(path_list):
    np_list = []
        
    for path in path_list:
        np_list.append(np.load(path))
           
    return np_list

with

def load_batch(path_list):
    return [np.load(path) for path in path_list]

In fact you can get rid of the function and directly use list comprehension. If functional call is required use lambda