HDF5 dataloaders using a python generator

[swpenninga]

Any dataloaders that use the H5Generator are effectively singlethreaded and cpu-bound:

zea/zea/data/dataloader.py

Line 376 in 2a3a9a6

def iterator(self):

Either I am completely misunderstanding the situation, but as far as I know:

tf.data.Dataset.from_generator(image_extractor, ...)

Means that next() runs in Python. This gives the pipeline:

GPU waits → Python → h5py → Python → TF → GPU

and num_workers, AUTOTUNE, prefetch, and batch have no effect in speeding this up.

Currently, I'm training an RF-data VAE that's taking 7s to load a batch, and 120ms to do the forward and backward pass which is why I cannot use the zea dataloader.

---

I understand that fixing this is quite a large task as there are many dependencies, but it would be good to look into at some point.

[tristan-deep]

That seems very inefficient indeed...

I guess one solution is to move from tf.data.Dataset.from_generator to tf.data.Dataset.from_tensor_slices?

[tristan-deep]

We can also look into other dataloaders that scale better, for instance grain. The gross of the logic from our dataloader is file and index handling, which can most likely be repurposed.

Also grain would work better with our preferred backend of choice (JAX), but does not require a JAX install, so it automatically is multi backend supportive.

[tristan-deep]

Working on the Grain solution in #256. From initial testing I think it would be good to transition to the Grain dataloader.

We would have to create some requirements / benchmarks so we actually build the dataloader that suits our needs.

[swpenninga]

Working on the Grain solution in #256. From initial testing I think it would be good to transition to the Grain dataloader.

We would have to create some requirements / benchmarks so we actually build the dataloader that suits our needs.

Maybe good to indeed make this list.
As a start, I assume we keep the whole zea format with hdf5 files for data storage?
Then what would be other requirements, simply speed? Benchmarking is relatively easy, as long as it is as fast as the current tensorflow loader.

[tristan-deep]

Agreed! Changing the hdf5 format would be quite a lot of work, and I'm not sure if there would be a great alternative that is as readable and flexible and also still memory efficient. Also, the current dataloader works for all hdf5 files, not necessarily zea format.

Besides speed, it would be nice to have some (but maybe not all) existing flexibility in place. Such as loading different axis or multiple frames etc.

Some requirement I can think off

• Speed of dataloader. Let's check with multiple loading settings / dataset sizes.
• File i/o. Can we index the hdf5 in different ways. There is probably an optimal way to store the data for fastest indexing, but one can always convert their dataset if they need that speed. In the past I found it super useful to use the same dataset, but batch over frames or other axes by specifying simply some arguments in the dataloader.
• Fancy dataloader options. Such as batching, caching, prefetching. I guess this ties into the speed thing mostly. Hopefully grain will implement that for us.

[swpenninga]

One thing I ran into is that if I wanted to add some non-standard data (like a latent-space encoding of RF data) I would need to add it under additional_elements and it would get put in the scan class.

I feel like there should perhaps be an entry like custom_data for stuff that does not fit raw_data, aligned_data, envelope_data, beamformed_data or image.

[tristan-deep]

As far as I know if you use the zea File or Dataset classes, you can only load properly zea formatted hdf5 files, indeed as you describe. The h5 dataloader however can currently load any hdf5 file if i'm not mistaken. You can specify any key.

Also additional_elements is the place to store custom_data (see example here), but maybe it is not the best name.