Могу ли я выполнить ленивое чтение с помощью h5py при нарезке с динамической осью?

У меня есть HDF5_generator, который возвращает такие данные:

for element_i in range(n_elements):
    img = f['data'][:].take(indices=element_i, axis=element_axis)
    yield img, label, weights

Я выполняю нарезку, потому что h5py, похоже, не предоставляет другого способа чтения (пожалуйста, поправьте меня, если я ошибаюсь), и я делаю это так (f['data'][:].take(...)), потому что хочу, чтобы ось среза быть динамичным и не знать, как делать "классические" нарезки (f['data'][:, :, element_i, :, :]) с динамической осью.

Но это ужасно медленно! Я даже не знаю, что происходит, потому что время чтения сильно колеблется, но я предполагаю, что для каждого element_i весь набор данных data считывается полностью, а иногда случайно все еще кэшируется, но иногда нет.

Я придумал «cache_full_file» (см. Полный код ниже), и это решает его:

cache_full_file = False
>> reading 19 elements/rows (poked data with (19, 4, 1024, 1024))
(4, 1024, 1024)  image read - Elapsed time: 6.5959 s            # every single read can take long
(4, 1024, 1024)  image read - Elapsed time: 28.0695 s
(4, 1024, 1024)  image read - Elapsed time: 0.6851 s
(4, 1024, 1024)  image read - Elapsed time: 3.3492 s
(4, 1024, 1024)  image read - Elapsed time: 0.5837 s
(4, 1024, 1024)  image read - Elapsed time: 1.0346 s
(4, 1024, 1024)  image read - Elapsed time: 2.5852 s
(4, 1024, 1024)  image read - Elapsed time: 18.7262 s
(4, 1024, 1024)  image read - Elapsed time: 19.1674 s           # ...


cache_full_file = True
>> reading 19 elements/rows (poked data with (19, 4, 1024, 1024))
(4, 1024, 1024)  image read - Elapsed time: 15.8334 s           # dataset is read and cached once
(4, 1024, 1024)  image read - Elapsed time: 0.0744 s            # following reads are all fast ...      
(4, 1024, 1024)  image read - Elapsed time: 0.0558 s            # ...

Но я не могу полагаться на полные файлы/наборы данных, помещающиеся в память!

Могу ли я выполнить "ленивое" чтение, при котором не считывается весь набор данных, чтобы извлечь фрагмент из набора данных HDF5?

Упрощенная версия кода класса:

class hdf5_generator:
    def __init__(self, file, repeat): self.file = file
    def __call__(self):
        with h5py.File(self.file, 'r') as f:
            n_elements = f['data'].shape[element_axis] # poke first dataset to get number of expected elements)

            if cache_full_file:
                img_eles = f['data'][:]     # read and store the whole dataset in memory
                for element_i in range(n_elements):
                    img = img_eles.take(indices=element_i, axis=element_axis)
                    yield img
            else:
                for element_i in range(n_elements):
                    # access a specific row in the dataset
                    img = f['data'][:].take(indices=element_i, axis=element_axis)
                    yield img, label, weights