Необъяснимое использование ОЗУ и потенциальная утечка памяти при использовании tf.data.TFRecordDataset

Question

Фон

Мы относительно новички в TensorFlow. Мы работаем над проблемой DL, связанной с набором видеоданных. Из-за большого объема данных мы решили предварительно обработать видео и сохранить кадры в формате jpeg в файлах TFRecord. Затем мы планируем использовать tf.data.TFRecordDataset для передачи данных в нашу модель.

Видео были обработаны в сегменты, каждый из которых состоит из 16 кадров, в сериализованном тензоре. Каждый кадр представляет собой изображение 128 * 128 RGB, которое было закодировано как jpeg. Каждый сериализованный сегмент сохраняется вместе с некоторыми метаданными в виде сериализованного tf.train.Example в TFRecords.

Версия TensorFlow: 2.1

Код

Ниже приведен код, который мы используем для создания tf.data.TFRecordDataset из TFRecords. Вы можете игнорировать поля num и file.

import os
import math
import tensorflow as tf

# Corresponding changes are to be made here
# if the feature description in tf2_preprocessing.py
# is changed
feature_description = {
    'segment': tf.io.FixedLenFeature([], tf.string),
    'file': tf.io.FixedLenFeature([], tf.string),
    'num': tf.io.FixedLenFeature([], tf.int64)
}


def build_dataset(dir_path, batch_size=16, file_buffer=500*1024*1024,
                  shuffle_buffer=1024, label=1):
    '''Return a tf.data.Dataset based on all TFRecords in dir_path
    Args:
    dir_path: path to directory containing the TFRecords
    batch_size: size of batch ie #training examples per element of the dataset
    file_buffer: for TFRecords, size in bytes
    shuffle_buffer: #examples to buffer while shuffling
    label: target label for the example
    '''
    # glob pattern for files
    file_pattern = os.path.join(dir_path, '*.tfrecord')
    # stores shuffled filenames
    file_ds = tf.data.Dataset.list_files(file_pattern)
    # read from multiple files in parallel
    ds = tf.data.TFRecordDataset(file_ds,
                                 num_parallel_reads=tf.data.experimental.AUTOTUNE,
                                 buffer_size=file_buffer)
    # randomly draw examples from the shuffle buffer
    ds = ds.shuffle(buffer_size=1024,
                    reshuffle_each_iteration=True)
    # batch the examples
    # dropping remainder for now, trouble when parsing - adding labels
    ds = ds.batch(batch_size, drop_remainder=True)
    # parse the records into the correct types
    ds = ds.map(lambda x: _my_parser(x, label, batch_size),
                num_parallel_calls=tf.data.experimental.AUTOTUNE)
    ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
    return ds


def _my_parser(examples, label, batch_size):
    '''Parses a batch of serialised tf.train.Example(s)
    Args:
    example: a batch serialised tf.train.Example(s)
    Returns:
    a tuple (segment, label)
    where segment is a tensor of shape (#in_batch, #frames, h, w, #channels)
    '''
    # ex will be a tensor of serialised tensors
    ex = tf.io.parse_example(examples, features=feature_description)
    ex['segment'] = tf.map_fn(lambda x: _parse_segment(x),
                              ex['segment'], dtype=tf.uint8)
    # ignoring filename and segment num for now
    # returns a tuple (tensor1, tensor2)
    # tensor1 is a batch of segments, tensor2 is the corresponding labels
    return (ex['segment'], tf.fill((batch_size, 1), label))


def _parse_segment(segment):
    '''Parses a segment and returns it as a tensor
    A segment is a serialised tensor of a number of encoded jpegs
    '''
    # now a tensor of encoded jpegs
    parsed = tf.io.parse_tensor(segment, out_type=tf.string)
    # now a tensor of shape (#frames, h, w, #channels)
    parsed = tf.map_fn(lambda y: tf.io.decode_jpeg(y), parsed, dtype=tf.uint8)
    return parsed

Проблема

Во время обучения наша модель вылетела из-за нехватки оперативной памяти. Мы исследовали, запустив несколько тестов и профилировав память с помощью memory-profiler с флагом --include-children.

Все эти тесты были запущены (только для ЦП) простым повторением набора данных несколько раз со следующим кодом:

count = 0
dir_path = 'some/path'
ds = build_dataset(dir_path, file_buffer=some_value)
for itr in range(100):
    print(itr)
    for itx in ds:
        count += 1

Общий размер подмножества TFRecords, над которым мы сейчас работаем, составляет ~ 3 ГБ. Мы бы предпочли использовать TF2.1, но мы также можем протестировать с TF2.2.

Согласно TF2 docs , file_buffer указывается в байтах.

Trial 1 : file_buffer = 500 *1024* 1024, TF2.1

Пробная версия 2 : file_buffer = 500 *1024* 1024, TF2.2 Этот вариант кажется намного лучше.

Пробная версия 3 file_buffer = 1024 * 1024, TF2.1 У нас нет графика, но максимальный объем ОЗУ составляет ~ 4.5 ГБ

Пробная версия 4 file_buffer = 1024 * 1024, TF2.1, но для предварительной выборки установлено значение 10

Я думаю, что здесь есть утечка памяти, поскольку мы видим использование памяти постепенно увеличивается.

Все испытания, указанные ниже, были выполнены только для 50 итераций вместо 100

Проба 5 file_buffer = 500 *1024* 1024, TF2.1, предварительная выборка = 2, все остальные значения AUTOTUNE были установлены на 16.

Пробная версия 6 file_buffer = 1024 * 1024, оставайтесь так же, как указано выше

Вопросы

1. Как значение file_buffer влияет на использование памяти, сравнение Trail 1 и Trail 3, file_buffer был уменьшен в 500 раз, но использование памяти уменьшилось только вдвое. Действительно ли значение файлового буфера выражено в байтах?
2. Параметры испытания 6 казались многообещающими, но попытка обучить модель с помощью того же метода не удалась, так как у него снова закончилась память.
3. Есть ли ошибка в TF2.1, почему огромная разница между испытанием 1 и испытанием 2?
4. Следует ли нам продолжать использовать AUTOTUNE или вернуться к постоянным значениям?

Я был бы рад провести больше тестов с разными параметрами. Заранее спасибо!