使用 Tensorflow Interleave 提高性能

Question

I have an input pipe that is performing poorly with low CPU, GPU, and disk utilization. I've been reading the tensorflow "Better performance with tf.data API" doc and the Dataset docs, but I don't understand what's going on well enough to apply it to my situation. Here's my current setup:

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)

Should I interleave the whole pipe right at the end (POINT2), before the prefetch? Or interleave imgd and outd separately, after each FixedLengthRecordDataset (POINT1A, POINT1B), and parallelize the maps? (need to keep the imgd and outd synced up!) What's up with Dataset.range(rvalue)---seems it's necessary but not obvious what rvalue to use? Is there a better overall plan?

Note that the datasets are very large and do not fit in RAM.

Answer 1

Interleave lets you process each file in a separate logical thread (in parallel), then combine the data from the files into a single dataset. Since your data comes from two corresponding files, you need to be careful to preserve the order.

Here is an example of how you could put the interleave near the end of the dataset:

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()

Each thread of the interleave will produce elements from a different pair of (img, out) files, and the elements produced from each pair of files will be interleaved together.