将 tf.data.Dataset.from_generator 与 TF2.1 并行化答案

【问题标题】：parallelising tf.data.Dataset.from_generator with TF2.1将 tf.data.Dataset.from_generator 与 TF2.1 并行化
【发布时间】：2020-02-07 17:14:58
【问题描述】：

他们已经有 2 篇关于此主题的帖子，但尚未针对最近的 TF2.1 版本进行更新...

简而言之，我有很多 tif 图像要使用特定管道读取和解析。

import tensorflow as tf
import numpy as np

files = # a list of str
labels = # a list of int
n_unique_label = len(np.unique(labels))

gen = functools.partial(generator, file_list=files, label_list=labels, param1=x1, param2=x2)
dataset = tf.data.Dataset.from_generator(gen, output_types=(tf.float32, tf.int32))
dataset = dataset.map(lambda b, c: (b, tf.one_hot(c, depth=n_unique_label)))

此处理效果很好。不过，我需要并行化文件解析部分，尝试以下解决方案：

files = # a list of str
files = tensorflow.data.Dataset.from_tensor_slices(files)

def wrapper(file_path):
    parser = partial(tif_parser, param1=x1, param2=x2)
    return tf.py_function(parser, inp=[file_path], Tout=[tf.float32])

dataset = files.map(wrapper, num_parallel_calls=2)

不同之处在于我在这里使用parser 函数一次解析一个文件。但是，它不起作用：

  File "loader.py", line 643, in tif_parser
    image = numpy.array(Image.open(file_path)).astype(float)

  File "python3.7/site-packages/PIL/Image.py", line 2815, in open
    fp = io.BytesIO(fp.read())

AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute 'read'


     [[{{node EagerPyFunc}}]] [Op:IteratorGetNextSync]

据我了解，tif_parser 函数接收的不是字符串而是（未计算的）张量。目前，这个函数相当简单：

def tif_parser(file_path, param1=1, param2=2):
    image = numpy.array(Image.open(file_path)).astype(float)
    image /= 255.0

    return image

【问题讨论】：

标签： python-3.x tensorflow tensorflow2.0 tensorflow-datasets

【解决方案1】：

我是这样处理的

dataset = tf.data.Dataset.from_tensor_slices((files, labels))

def wrapper(file_path, label):
    import functools
    parser = functools.partial(tif_parser,  param1=x1, param2=x2)
    return tf.data.Dataset.from_generator(parser, (tf.float32, tf.int32), args=(file_path, label))

dataset = dataset.interleave(wrapper, cycle_length=tf.data.experimental.AUTOTUNE)

# The labels are converted to 1-hot vectors, could be integrated in tif_parser
dataset = dataset.map(lambda i, l: (i, tf.one_hot(l, depth=unique_label_count)))

dataset = dataset.shuffle(buffer_size=file_count, reshuffle_each_iteration=True)
dataset = dataset.batch(batch_size=batch_size, drop_remainder=False)
dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

具体来说，每次调用解析器时，我都会生成一个数据集。解析器在每次调用时运行cycle_length 时间，这意味着一次读取cycle_length 图像。这适合我的具体情况，因为我无法将所有图像加载到内存中。我不确定预取是否在这里正确使用。

【讨论】：