6

背景

我们对 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

所有这些测试都是通过使用以下代码多次迭代数据集来运行的（仅限 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 文档，file_buffer 以字节为单位。

试用一：file_buffer = 500*1024*1024，TF2.1 在此处输入图像描述

试用二：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, prefetch = 2，所有其他 AUTOTUNE 值都设置为 16。 在此处输入图像描述

试用6 file_buffer = 1024*1024，其余同上 在此处输入图像描述

问题

  1. file_buffer 值对内存使用有什么影响，对比 Trail 1 和 Trail 3，file_buffer 减少了 500 倍，但内存使用量只减少了一半。文件缓冲区值真的以字节为单位吗？
  2. 试验 6 的参数看起来很有希望，但尝试用相同的参数训练模型失败了，因为它再次耗尽了内存。
  3. TF2.1有bug吗，为什么试玩1和试玩2差别这么大？
  4. 我们应该继续使用 AUTOTUNE 还是恢复到恒定值？

我很乐意使用不同的参数运行更多测试。提前致谢！

4

0 回答 0