Python tf.data.TFRecordDataset.interleave用法及代码示例

用法

interleave(
    map_func, cycle_length=None, block_length=None, num_parallel_calls=None,
    deterministic=None, name=None
)

跨此数据集映射map_func，并交错结果。

类型签名是：

def interleave(
  self:Dataset[T],
  map_func:Callable[[T], Dataset[S]]
) -> Dataset[S]

例如，您可以使用 Dataset.interleave() 来同时处理多个输入文件：

# Preprocess 4 files concurrently, and interleave blocks of 16 records
# from each file.
filenames = ["/var/data/file1.txt", "/var/data/file2.txt",
             "/var/data/file3.txt", "/var/data/file4.txt"]
dataset = tf.data.Dataset.from_tensor_slices(filenames)
def parse_fn(filename):
  return tf.data.Dataset.range(10)
dataset = dataset.interleave(lambda x:
    tf.data.TextLineDataset(x).map(parse_fn, num_parallel_calls=1),
    cycle_length=4, block_length=16)

cycle_length 和 block_length 参数控制生成元素的顺序。 cycle_length 控制同时处理的输入元素的数量。如果将 cycle_length 设置为 1，则此转换将一次处理一个输入元素，并将产生与 tf.data.Dataset.flat_map 相同的结果。通常，此转换会将map_func 应用于cycle_length 输入元素，在返回的Dataset 对象上打开迭代器，并在它们之间循环，从每个迭代器生成block_length 连续元素，并在每次执行时消耗下一个输入元素到达迭代器的末尾。

例如：

dataset = Dataset.range(1, 6)  # ==> [ 1, 2, 3, 4, 5 ]
# NOTE:New lines indicate "block" boundaries.
dataset = dataset.interleave(
    lambda x:Dataset.from_tensors(x).repeat(6),
    cycle_length=2, block_length=4)
list(dataset.as_numpy_iterator())
[1, 1, 1, 1,
 2, 2, 2, 2,
 1, 1,
 2, 2,
 3, 3, 3, 3,
 4, 4, 4, 4,
 3, 3,
 4, 4,
 5, 5, 5, 5,
 5, 5]

注意：只要 map_func 是纯函数且 deterministic=True ，此转换产生的元素的顺序是确定的。如果map_func 包含任何有状态操作，则访问该状态的顺序是未定义的。

性能通常可以通过设置num_parallel_calls 来提高，这样interleave 将使用多个线程来获取元素。如果不需要确定性，设置 deterministic=False 也可以提高性能。

filenames = ["/var/data/file1.txt", "/var/data/file2.txt",
             "/var/data/file3.txt", "/var/data/file4.txt"]
dataset = tf.data.Dataset.from_tensor_slices(filenames)
dataset = dataset.interleave(lambda x:tf.data.TFRecordDataset(x),
    cycle_length=4, num_parallel_calls=tf.data.AUTOTUNE,
    deterministic=False)