tensorflow - 获取文件名示例来自 tf,parse_example

Question

我正在 tensorflow 中编写一个数据输入管道，它使用一堆具有不同示例(类型)的 tfrecord 文件。

我正在使用如下代码:

filenames = ["/var/data/file1.tfrecord", "/var/data/file2.tfrecord"]
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.map(_parse_function)

但是我希望我的 parse_function 对于 file1.tfrecord 与 file2.tfrecord 不同。我如何实现这一点。在 parse_example 中是否有某种方式知道特定示例来自哪个文件？

Answer 1

您可以使用 Dataset.flat_map()将文件名包含在每条记录中的转换如下:

filenames = ["/var/data/file1.tfrecord", "/var/data/file2.tfrecord"]
filenames = tf.data.from_tensor_slices(filenames)

# `Dataset.flat_map()` creates a nested dataset from each element in `filenames`.
#
# For each file in filename, zip together the filename (repeated infinitely) with
# the records read from that file.
dataset = filenames.flat_map(
    lambda fn: tf.data.Dataset.zip((tf.data.Dataset.from_tensors(fn).repeat(None),
                                    tf.data.TFRecordDataset(fn))))

# The _parse_function can now be modified to take both the filename and the record.
dataset = dataset.map(lambda fn, record: _parse_function(fn, record))