Loading speed vs memory: how to efficiently load large arrays from h5 file

Viewed 719

I have been facing the following issue: I have to loop over num_objects = 897 objects, for every single one of which I have to use num_files = 2120 h5 files. The files are very large, each being 1.48 GB and the content of interest to me is 3 arrays of floats of size 256 x 256 x 256 contained in every file (v1,v2 and v3). This is to say, the loop looks like:

for i in range(num_objects):
    ...
    for j in range(num_files):
       some operation with the three 256 x 256 x 256 arrays in each file

and my current way of loading them is by doing the following in the innermost loop:

f = h5py.File('output_'+str(q)+'.h5','r')
key1 = np.array(f['key1'])
v1=key1[:,:,:,0]
v2=key2[:,:,:,1]
v3=key3[:,:,:,2]

The above option of loading the files every time for every object is obviously very slow. On the other hand, loading all files at once and importing them in a dictionary results in excessive use of memory and my job gets killed. Some diagnostics:

  • The above way takes 0.48 seconds per file, per object, resulting in a total of 10.5 days (!) spent only on this operation.
  • I tried exporting key1 to npz files, but it was actually slower by 0.7 seconds per file.
  • I exported v1,v2 and v3 individually for every file to npz files (i.e. 3 npz files for each h5 file), but that saved me only 1.5 day in total.

Does anyone have some other idea/suggestion I could try to be fast and at the same time not limited by excessive memory usage?

3 Answers

If I understand, you have 2120 .h5 files. Do you only read the 3 arrays in dataset f['key1'] for each file? (or are there other datasets?) If you only/always read f['key1'], that's a bottleneck you can't program around. Using a SSD will help (because I/O is faster than HDD). Otherwise, you will have to reorganize your data. The amount of RAM on your system will determine the number of arrays you can read simultaneously. How much RAM do you have?

You might gain a little speed with a small code change. v1=key1[:,:,:,0] returns v1 as an array (same for v2 and v3). There's no need to read dataset f['key1'] into an array. Doing that doubles your memory footprint. (BTW, Is there a reason to convert your arrays to a dictionary?)

The process below only creates 3 arrays by slicing v1,v2,v3 from the h5py f['key1']object. It will reduce your memory footprint for each loop by 50%.

f = h5py.File('output_'+str(q)+'.h5','r')
key1 = f['key1'] 
## key1 is returned as a h5py dataset OBJECT, not an array
v1=key1[:,:,:,0]
v2=key2[:,:,:,1]
v3=key3[:,:,:,2]

On the HDF5 side, since you always slice out the last axis, your chunk parameters might improve I/O. However, if you want to change the chunk shape, you will have to recreate the .h5 files. So, that will probably not save time (at least in the short-term).

Agreed with @kcw78, in my experience the bottleneck is usually that you load too much data compared to what you need. So don't convert the dataset to an array unless you need to, and slice only the part you need (do you need the whole [:,:,0] or just a subset of it?).

Also, if you can, change the order of the loops, so that you open each file only once.

for j in range(num_files):
    ...
    for i in range(num_objects):
       some operation with the three 256 x 256 x 256 arrays in each file

Another approach would be to use an external tool (like h5copy) to extract the data you need in a much smaller dataset, and then read it in python, to avoid python's overhead (which might not be that much to be honest).

Finally, you could use multiprocessing to take advantage of you CPU cores, as long as you tasks are relatively independent from one another. You have an example here.

h5py is the data representation class the supports most of NumPy data types and supports traditional NumPy operations like slicing, and a variety of descriptive attributes such as shape and size attributes.

There is a cool option that you can use for representing datasets in chunks using the create_dataset method with the chunked keyword enabled:

dset = f.create_dataset("chunked", (1000, 1000), chunks=(100, 100))

The benefit of this is that you can easily resize your dataset and now as in your case, you would only have to read a chunk and not the whole 1.4 GB of data.

But beware of chunking implications: different data sizes work best of different chunk sizes. There is an option of autochunk that selects this chunk size automatically rather than through hit and trial:

dset = f.create_dataset("autochunk", (1000, 1000), chunks=True)

Cheers

Related