使用Tensorflow交织以提高性能

Question

我的输入管道在CPU，GPU和磁盘使用率较低的情况下表现不佳。我一直在阅读tensorflow“使用tf.data API更好的性能”文档和Dataset文档，但我不了解发生了什么事情足以将其应用于我的情况。这是我当前的设置：

img_files = sorted(tf.io.gfile.glob(...))
imgd = tf.data.FixedLengthRecordDataset(img_files, inrez*inrez)
#POINT1A
imgd = imgd.map(lambda s: tf.reshape(tf.io.decode_raw(s, tf.int8), (inrez,inrez,1)))
imgd = imgd.map(lambda x: tf.cast(x, dtype=tf.float32))

out_files = sorted(tf.io.gfile.glob(...))
outd = tf.data.FixedLengthRecordDataset(out_files, 4, compression_type="GZIP")
#POINT1B
outd = outd.map(lambda s: tf.io.decode_raw(s, tf.float32))

xsrc = tf.data.Dataset.zip((imgd, outd)).batch(batchsize)
xsrc = xsrc.repeat()        # indefinitely
#POINT2
xsrc = xsrc.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

我应该在预取之前在末端（POINT2）插入整个管道吗？还是在每个FixedLengthRecordDataset（POINT1A，POINT1B）之后分别交织imgd和outd，并并行化地图？ （需要使imgd和outd保持同步！）Dataset.range（rvalue）是怎么回事-似乎有必要但不明显要使用什么rvalue？有没有更好的总体计划？

请注意，数据集非常大，不适合放入RAM。

Answer 1

Interleave使您可以在单独的逻辑线程中（并行）处理每个文件，然后将文件中的数据组合到单个数据集中。由于您的数据来自两个相应的文件，因此您需要注意保留顺序。

这是一个如何将交织放置在数据集结尾附近的示例：

img_files = ...
out_files = ...
files = tf.data.Dataset.zip(img_files, out_files)

def parse_img_file(img_file):
  imgd = tf.data.FixedLengthRecordDataset(img_files, inrez*inrez)
  ...

def parse_out_file(out_file):
  ...

def parse_files_fn(img_file, out_file):
  img_file_dataset = parse_img_file(img_file)
  out_file_dataset = parse_out_file(out_file)
  return tf.data.Dataset.zip(img_file_dataset, out_file_dataset)

dataset = files.interleave(parse_files_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.repeat()

交错的每个线程将从不同的一对（img，out）文件中产生元素，而从每一对文件中产生的元素将被交错在一起。