如何在 TF 2 中使用带有自定义函数的 tf.data.Dataset.interleave()?

Question

我正在使用TF 2.2，我正在尝试使用tf.data来创建管道。

以下工作正常：

def load_image(filePath, label):
print('Loading File: {}' + filePath)
raw_bytes = tf.io.read_file(filePath)
image = tf.io.decode_image(raw_bytes, expand_animations = False)
return image, label
# TrainDS Pipeline
trainDS = getDataset()
trainDS = trainDS.shuffle(size['train'])
trainDS = trainDS.map(load_image, num_parallel_calls=AUTOTUNE)
for d in trainDS:
print('Image: {} - Label: {}'.format(d[0], d[1]))

我想将load_image()与Dataset.interleave()一起使用.然后我尝试了：

# TrainDS Pipeline
trainDS = getDataset()
trainDS = trainDS.shuffle(size['train'])
trainDS = trainDS.interleave(lambda x, y: load_image_with_label(x, y), cycle_length=4)
for d in trainDS:
print('Image: {} - Label: {}'.format(d[0], d[1]))

但是我收到以下错误：

Exception has occurred: TypeError
`map_func` must return a `Dataset` object. Got <class 'tuple'>
File "/data/dev/train_daninhas.py", line 44, in <module>
trainDS = trainDS.interleave(lambda x, y: load_image_with_label(x, y), cycle_length=4)

如何调整我的代码以使Dataset.interleave()与load_image()并行读取图像？

Answer 1

正如错误所暗示的那样，您需要修改load_image以使其返回Dataset对象，我展示了一个示例，其中包含两个图像，介绍如何在tensorflow 2.2.0中执行此操作：

import tensorflow as tf
filenames = ["./img1.jpg", "./img2.jpg"]
labels = ["A", "B"]
def load_image(filePath, label):
print('Loading File: {}' + filePath)
raw_bytes = tf.io.read_file(filePath)
image = tf.io.decode_image(raw_bytes, expand_animations = False)
return tf.data.Dataset.from_tensors((image, label))
dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
dataset = dataset.interleave(lambda x, y: load_image(x, y), cycle_length=4)
for i in dataset.as_numpy_iterator():
image = i[0]
label = i[1]
print(image.shape)
print(label.decode())
# (275, 183, 3)
# A
# (275, 183, 3)
# B