Блог

Я хочу создать конвейер данных для обучения моделей TensorFlow. Данные хранятся в очень больших файлах HDF5 (250 ГБ).

Я написал конвейер, который работает для входных файлов меньшего размера, но в конечном итоге уничтожается ядром после использования слишком большого объема оперативной памяти подкачки (проверено с помощью мониторинга).

 import tensorflow as tf
import h5py

class TestGenerator:
    """
    Implements a generator that can be used by tf.data.Dataset.from_generator 
    to produce a dataset for any test data.
    """
    def __init__(self, src, dset):
        self.src = src
        self.dset = dset
        self.output_signature = (
            tf.TensorSpec(shape=(2,), dtype=tf.uint64)
        )

    def __call__(self):
        """This is needed for tf.data.Dataset.from_generator to work."""
        with h5py.File(self.src, 'r', swmr=True) as f:
            for sample in f[self.dset]:
                yield sample[0], sample[1]

gen = TestGenerator('h5file.h5', 'dset_path')

dataset = tf.data.Dataset.from_generator(
    gen, 
    output_signature=gen.output_signature
)

for sample in dataset:
    pass
 

Сначала я подумал, что это может быть проблемой модуля h5py, поэтому я протестировал его отдельно:

 with h5py.File('h5file.h5', 'r', swmr=True) as f:
    for sample in f['dset_path']:
        pass
 

Это работает без проблем. Это приводит к выводу, что TensorFlow ответственен за проблему с памятью. Что меня раздражает, так это то, что я предположил, что TensorFlow извлекает необходимые данные «на лету» и, следовательно, позволяет избежать проблем с памятью.

Код протестирован и отлично работает для файлов меньшего размера. Я также тестировал версии, которые я использовал dataset.prefetch перед итерацией, но с тем же результатом.

Загружает ли TensorFlow весь набор данных под капотом?

Если вы откроете диспетчер задач во время импорта tensorflow и создания вашей модели, покажет ли он чрезвычайно высокие значения зарезервированной памяти GPU?

Если это так, то это может быть не размер модели или данных, а тот факт, что tensorflow резервирует как можно больше памяти для всего обучения модели.

Для моего личного примера, мой 3080 имеет 10 ГБ выделенной памяти GPU, а tensorflow занял 9,7 ГБ.

Если это так, посмотрите на метод set_memory_growth на https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth .

Использование этого уменьшило использование выделенной памяти GPU с 9,7 ГБ до 3,2 ГБ и 4 ГБ.

РЕДАКТИРОВАТЬ: я не уверен, в какой момент он будет стремиться выделить эту память, но если вы попытаетесь обучить модель и следить за производительностью диспетчера задач для графического процессора, вы должны выяснить, ведет ли он себя таким образом!