如何提高数据输入管道的性能？

Question

我尝试优化我的数据输入管道。数据集是一组 450 个 TFRecord 文件，每个文件大小约为 70MB，托管在 GCS 上。该作业使用 GCP ML Engine 执行。没有GPU。

这是管道：

def build_dataset(file_pattern):
    return tf.data.Dataset.list_files(
        file_pattern
    ).interleave(
        tf.data.TFRecordDataset,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    ).shuffle(
        buffer_size=2048
    ).batch(
        batch_size=2048,
        drop_remainder=True,
    ).cache(
    ).repeat(
    ).map(
        map_func=_parse_example_batch,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    ).prefetch(
        buffer_size=1
    )

使用映射函数：

def _bit_to_float(string_batch: tf.Tensor):
    return tf.reshape(tf.math.floormod(tf.dtypes.cast(tf.bitwise.right_shift(
        tf.expand_dims(tf.io.decode_raw(string_batch, tf.uint8), 2),
        tf.reshape(tf.dtypes.cast(tf.range(7, -1, -1), tf.uint8), (1, 1, 8))
    ), tf.float32), 2), (tf.shape(string_batch)[0], -1))


def _parse_example_batch(example_batch):
    preprocessed_sample_columns = {
        "features": tf.io.VarLenFeature(tf.float32),
        "booleanFeatures": tf.io.FixedLenFeature((), tf.string, ""),
        "label": tf.io.FixedLenFeature((), tf.float32, -1)
    }
    samples = tf.io.parse_example(example_batch, preprocessed_sample_columns)
    dense_float = tf.sparse.to_dense(samples["features"])
    bits_to_float = _bit_to_float(samples["booleanFeatures"])
    return (
        tf.concat([dense_float, bits_to_float], 1),
        tf.reshape(samples["label"], (-1, 1))
    )

我尝试遵循 data pipeline tutorial, and vectorize my mapped function (as advised by ) 的最佳实践。

使用此设置，在高速下载数据时（带宽约为 200MB/s），CPU 使用不足 (14%) 并且训练速度非常慢（超过 1 小时一个时代）。

我尝试了一些参数配置，更改了 interleave() 参数，如 num_parallel_calls 或 cycle_length 或 TFRecordDataset 参数，如 num_parallel_calls.

最快的配置使用这组参数：

interleave.num_parallel_calls: 1
interleave.cycle_length: 8
TFRecordDataset.num_parallel_calls: 8

有了这个，一个 epoch 只需约 20 分钟即可到达运行。 但是，CPU使用率仅为 50%，而带宽消耗约为 55MB/s

问题：

如何优化管道以达到 100% CPU 使用率（以及大约 100MB/s 的带宽消耗）？
为什么 tf.data.experimental.AUTOTUNE 找不到加速训练的最佳值？

善良，亚历克西斯

编辑

经过更多的实验，我得出了以下解决方案。

如果 num_parallel_calls 大于 0，则删除 TFRecordDataset 已经处理的 interleave 步骤。
将映射函数更新为仅执行 parse_example 和 decode_raw，返回一个元组 `((, ), ())
cache后map
将 _bit_to_float 函数移动为模型的一个组件

最后，这里是数据管道代码：

def build_dataset(file_pattern):
    return tf.data.TFRecordDataset(
        tf.data.Dataset.list_files(file_pattern),
        num_parallel_reads=multiprocessing.cpu_count(),
        buffer_size=70*1000*1000
    ).shuffle(
        buffer_size=2048
    ).map(
        map_func=split,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    ).batch(
        batch_size=2048,
        drop_remainder=True,
    ).cache(
    ).repeat(
    ).prefetch(
        buffer_size=32
    )


def split(example):
    preprocessed_sample_columns = {
        "features": tf.io.VarLenFeature(tf.float32),
        "booleanFeatures": tf.io.FixedLenFeature((), tf.string, ""),
        "label": tf.io.FixedLenFeature((), tf.float32, -1)
    }
    samples = tf.io.parse_single_example(example, preprocessed_sample_columns)
    dense_float = tf.sparse.to_dense(samples["features"])
    bits_to_float = tf.io.decode_raw(samples["booleanFeatures"], tf.uint8)
    return (
        (dense_float, bits_to_float),
        tf.reshape(samples["label"], (1,))
    )


def build_model(input_shape):
    feature = keras.Input(shape=(N,))
    bool_feature = keras.Input(shape=(M,), dtype="uint8")
    one_hot = dataset._bit_to_float(bool_feature)
    dense_input = tf.reshape(
        keras.backend.concatenate([feature, one_hot], 1),
        input_shape)
    output = actual_model(dense_input)

    model = keras.Model([feature, bool_feature], output)
    return model

def _bit_to_float(string_batch: tf.Tensor):
    return tf.dtypes.cast(tf.reshape(
        tf.bitwise.bitwise_and(
            tf.bitwise.right_shift(
                tf.expand_dims(string_batch, 2),
                tf.reshape(
                    tf.dtypes.cast(tf.range(7, -1, -1), tf.uint8),
                    (1, 1, 8)
                ),
            ),
            tf.constant(0x01, dtype=tf.uint8)
        ),
        (tf.shape(string_batch)[0], -1)
    ), tf.float32)

感谢所有这些优化：

带宽消耗约为 90MB/s
CPU 使用率约为 20%
第一个纪元花费 20 分钟
连续的 epoch 每个花费 5 分钟

所以这似乎是一个很好的初始设置。但是CPU和BW还没有被过度使用，所以仍然欢迎任何建议！

编辑二

因此，经过一些基准测试后，我发现了我认为最好的输入管道：

def build_dataset(file_pattern):
    tf.data.Dataset.list_files(
        file_pattern
    ).interleave(
        TFRecordDataset,
        cycle_length=tf.data.experimental.AUTOTUNE,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    ).shuffle(
        2048
    ).batch(
        batch_size=64,
        drop_remainder=True,
    ).map(
        map_func=parse_examples_batch,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    ).cache(
    ).prefetch(
        tf.data.experimental.AUTOTUNE
    )

def parse_examples_batch(examples):
    preprocessed_sample_columns = {
        "features": tf.io.FixedLenSequenceFeature((), tf.float32, allow_missing=True),
        "booleanFeatures": tf.io.FixedLenFeature((), tf.string, ""),
        "label": tf.io.FixedLenFeature((), tf.float32, -1)
    }
    samples = tf.io.parse_example(examples, preprocessed_sample_columns)
    bits_to_float = tf.io.decode_raw(samples["booleanFeatures"], tf.uint8)
    return (
        (samples['features'], bits_to_float),
        tf.expand_dims(samples["label"], 1)
    )

所以，有什么新鲜事：

根据这个GitHub issue，TFRecordDataset交织是一个遗留的，所以interleave功能更好。
batch在map之前是个好习惯（vectorizing your function），减少映射函数的调用次数。
不再需要 repeat。从TF2.0开始，Keras模型API支持数据集API，可以使用缓存（见）
从 VarLenFeature 切换到 FixedLenSequenceFeature，删除对 tf.sparse.to_dense 的无用调用。

希望这对您有所帮助。仍然欢迎建议。

Answer 1

为了社区的利益，在答案部分提及@AlexisBRENON 的解决方案和重要观察。

以下是重要观察结果：

根据这个GitHub issue，TFRecordDataset interleaving 是遗留的，所以interleave 功能更好。
batch before map 是一个好习惯（vectorizing your function），减少调用映射函数的次数。
不再需要 repeat。从 TF2.0 开始，Keras 模型 API 支持数据集 API 并且可以使用缓存（参见）
从 VarLenFeature 切换到 FixedLenSequenceFeature，删除对 tf.sparse.to_dense 的无用调用。

管道代码，具有改进的性能，符合上述观察结果如下：

def build_dataset(file_pattern):
    tf.data.Dataset.list_files(
        file_pattern
    ).interleave(
        TFRecordDataset,
        cycle_length=tf.data.experimental.AUTOTUNE,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    ).shuffle(
        2048
    ).batch(
        batch_size=64,
        drop_remainder=True,
    ).map(
        map_func=parse_examples_batch,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    ).cache(
    ).prefetch(
        tf.data.experimental.AUTOTUNE
    )

def parse_examples_batch(examples):
    preprocessed_sample_columns = {
        "features": tf.io.FixedLenSequenceFeature((), tf.float32, allow_missing=True),
        "booleanFeatures": tf.io.FixedLenFeature((), tf.string, ""),
        "label": tf.io.FixedLenFeature((), tf.float32, -1)
    }
    samples = tf.io.parse_example(examples, preprocessed_sample_columns)
    bits_to_float = tf.io.decode_raw(samples["booleanFeatures"], tf.uint8)
    return (
        (samples['features'], bits_to_float),
        tf.expand_dims(samples["label"], 1)
    )

Answer 2

我还有一个建议要补充：

根据interleave()的文档，你可以作为第一个参数使用映射函数。

这意味着，可以这样写：

 dataset = tf.data.Dataset.list_files(file_pattern)
 dataset = dataset.interleave(lambda x:
    tf.data.TFRecordDataset(x).map(parse_fn, num_parallel_calls=AUTOTUNE),
    cycle_length=tf.data.experimental.AUTOTUNE,
    num_parallel_calls=tf.data.experimental.AUTOTUNE
    )

据我了解，这会将一个解析函数映射到每个分片，然后交错结果。这样就消除了以后 dataset.map(...) 的使用。

如何提高数据输入管道的性能？

How to improve data input pipeline performance?

python

python-3.x

tensorflow

tensorflow-datasets

tensorflow2.0

问题：

编辑

编辑二