使用tf.data.TFRecordDataset

Question

背景

我们是TensorFlow的新手。我们正在处理涉及视频数据集的DL问题。由于涉及的数据量大，我们决定对视频进行预处理，并将帧以jpeg格式存储在TFRecord文件中。然后，我们计划使用tf.data.TFRecordDataset将数据提供给我们的模型。

视频已按序列化张量处理为分段，每个分段包含16帧。每帧是一个128 * 128 RGB图像，被编码为jpeg。每个序列化段都与一些元数据一起存储为TFRecords中的序列化tf.train.Example。

TensorFlow版本：2.1

代码

下面是我们用来从TFRecords创建tf.data.TFRecordDataset的代码。您可以忽略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

问题

在训练时，我们的模型因RAM不足而崩溃。我们通过运行一些测试并对带有--include-children标志的memory-profiler进行了内存分析来进行调查。

通过使用以下代码简单地多次遍历数据集来运行所有这些测试（仅限CPU）：

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子集的总大小约为3GB 我们宁愿使用TF2.1，但也可以使用TF2.2进行测试。

根据TF2 docs，file_buffer以字节为单位。

试验1 ：file_buffer = 500 * 1024 * 1024，TF2.1

试用2 ：file_buffer = 500 * 1024 * 1024，TF2.2 这个似乎好多了。

试用3 file_buffer = 1024 * 1024，TF2.1 我们没有这块土地，但RAM最高可达〜4.5GB

试验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中是否存在错误，为什么Trial 1和Trial 2之间存在巨大差异？
4. 我们应该继续使用AUTOTUNE还是恢复为常数值？

我很乐意运行更多使用不同参数的测试。 预先感谢！