Repository: tensorflow
Issue: Memory usage with tf.data pipeline (HDF5, TFRecords)
Labels: stat:awaiting tensorflower,comp:data,type:performance,2.17

### Issue type

Performance

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

tf 2.17.0

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 22.04.4 LTS

### Mobile device

_No response_

### Python version

3.11.9

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

12.3

### GPU model and memory

RTX A6000 Ada

### Current behavior?

Hello.

I am wondering about memory usage for the Tensorflow data API. I am running into OOM issues when I run on very large (360,000 files, >40 MB per file) files. I cannot even iterate through the dataset appropriately; forget about training a (Keras) model.

Here is memory usage (using `psutil` and `memory_info().rss`) over 5 "epochs" where I iterate over a dummy HDF5 dataset. The drops that create the sawtooth pattern are each "epoch," where I just iterate over the dataset 5 times.

![image](https://github.com/user-attachments/assets/93e2f201-9c3a-411f-83c6-5861d378fdd1)

Is this expected behavior? I can repeat this memory curve with `TFRecordDataset` with and without `interleave`.

[Gist for creating dummy hdf5 files](https://gist.github.com/dryglicki/dec05492b73416e3829a0440c6024793)
[Gist for reading dummy hdf5 files with Dataset API](https://gist.github.com/dryglicki/cf0a52dd31af3358d5e3cd5351c51f13)

Be warned, that to get the plotting to work without wrecking Tensorflow, you need to use matplotlib <3.8, since that is when it incorporates Numpy 2.0 -- I used MPL v3.7.3.

### Standalone code to reproduce the issue

```shell
#!/usr/bin/env python3

import os ; import sys
from pathlib import Path
import glob

import tensorflow as tf

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt

import h5py
import psutil
import gc


def unpack_hdf5(hdf5_file):
    '''
    Simple and quick unpack of an hdf5 dummy file for tensorflow.
    '''
    with h5py.File(hdf5_file.numpy().decode('utf-8'), 'r') as h5:
        priors = tf.convert_to_tensor(h5.get('priors')[...], tf.float32)
        forecast = tf.convert_to_tensor(h5.get('forecast')[...], tf.float32)
        model = tf.convert_to_tensor(h5.get('model')[...], tf.float32)

    return priors, model, forecast

def interleave_wrapper(hdf5_file,
        prior_shape,
        model_shape,
        forecast_shape):
    '''
    Wrapper for py_function to work with interleave. Interleave requires a dataset upon output.
    '''
    p, m, f = tf.py_function(unpack_hdf5, [hdf5_file],
                             Tout = [tf.float32, tf.float32, tf.float32])

    p.set_shape(prior_shape)
    m.set_shape(model_shape)
    f.set_shape(forecast_shape)

    return tf.data.Dataset.zip( (tf.data.Dataset.from_tensors((p,m)).map(
        lambda p, m: {'priors' : p, 'model' : m}, num_parallel_calls = tf.data.AUTOTUNE),
        tf.data.Dataset.from_tensors(f)) )

def create_dataset_interleave(file_list,
        priors_shape,
        model_shape,
        forecast_shape,
        batch_size = 32):
    '''
    Create a dataset using pre-globbed file list
    '''

    return tf.data.Dataset.from_tensor_slices(file_list).interleave(
            lambda x: interleave_wrapper(x, priors_shape, model_shape, forecast_shape),
            cycle_length = tf.data.AUTOTUNE,
            num_parallel_calls = tf.data.AUTOTUNE,
            deterministic = False).batch(
                    batch_size, drop_remainder=True)

def main():
    '''
    Program to monitor memory usage of HDF5 file reader.
    '''
    p = psutil.Process(os.getpid())
    hdf_directory = 'hdf_files'
    file_list = list(glob.glob(f'{hdf_directory}/*hdf5'))
    print(len(file_list))

    test_file = file_list[0]
    hdf_vars = ['priors', 'model', 'forecast']
    shapes = {}
    with h5py.File(test_file, 'r') as h5:
        for var in hdf_vars:
            dict_var = f'{var}_shape'
            shapes[dict_var] = h5[var].shape

    ds = create_dataset_interleave(file_list, **shapes)

    rssUse = []
    batch_numbers = []
    for epoch in range(5):
        for ii, batch in enumerate(ds):
            batch_numbers.append((epoch+1)*(ii+1))
            X, Y = batch
            print('Batch: ',ii)
            print(X['model'].shape)
            rssUse.append(p.memory_info().rss / (1024 ** 2))
        gc.collect()

    for b, r in zip(batch_numbers, rssUse):
        print(f'{b:04d}: {r}')

    fig = plt.figure(figsize = (4,3))
    ax = fig.add_subplot()

    ax.plot(rssUse, label = 'rss')
    ax.legend()
    ax.set_ylabel('Memory usage [MB]')
    ax.set_xlabel('Step')

    fig.savefig('memory_use_hdf5.png', dpi=200, bbox_inches='tight')


if __name__ == '__main__':
    main()
```


### Relevant log output

_No response_