使用 Tensorflow Interleave 提高性能

Question

我有一個輸入管道，它在 CPU、GPU 和磁盤利用率低的情況下表現不佳。 我一直在閱讀 tensorflow “Better performance with tf.data API” doc 和 Dataset docs，但我不明白發生了什么足以將其應用於我的情況。 這是我目前的設置：

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) 怎么了---似乎有必要但不明顯使用什么右值？ 有沒有更好的整體方案？

請注意，數據集非常大，不適合 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) 文件對生成元素，並且從每對文件生成的元素將被交錯在一起。