TensorFlow py_function 嵌套输出类型？

Question

是否可以为 TensorFlows 的 py_function 指定嵌套输出类型？

作为一个具体案例，我希望 py_function 具有 return 类型的 ((tf.float32, tf.float32), (tf.float32, tf.float32))，其中各个元素不一定具有相同的维度。有没有办法为 py_function?

指定这个

正如了解为什么这对我的情况有用一样，我有一个 tf.data.Dataset 和文件路径列表。 py_function 采用这些文件路径之一，并从文件中生成一个负面和正面的例子以及相应的标签，从而产生 ((positive_data, positive_label), (negative_data, negative_label)) （注意，标签不一定是单个值，但它们也是与输入数据的形状不同）。这个 py_function 可以映射到数据集，并且（具有上述结构）将一个级别展平以生成具有 (data, label) 结构化元素的训练数据集。虽然可以有一种解决方法，即数据和标签堆叠在 py_function 中，然后再取消堆叠（或者从 py_function 开始完全非结构化，然后才配对），但这会导致混乱和混乱的设置.如果 py_function 可以直接输出 ((tf.float32, tf.float32), (tf.float32, tf.float32)) 类型，它会导致更清晰的设置。

Answer 1

tf.py_function的输出类型不能是嵌套序列。然而，当 tf.py_function 与 tf.data API 一起使用时，您需要创建一个包装函数（下例中的 tf_foo），并且您可以将输出嵌套在其中函数。

import tensorflow as tf

# The python function.
def foo(x):
    return x, x, x, x

# Wrap the python function to make it compatible with `tf.data.Dataset.map`.
def tf_foo(x):
    a, b, c, d = tf.py_function(foo, [x], Tout=[tf.float32, tf.float32, tf.float32, tf.float32])
    return (a, b), (c, d)

dset = tf.data.Dataset.from_tensor_slices([0, 1, 2, 3, 4])
dset.map(tf_foo)
# <MapDataset shapes: ((<unknown>, <unknown>), (<unknown>, <unknown>)),
#  types: ((tf.float32, tf.float32), (tf.float32, tf.float32))>

这在TensorFlow tutorial中也有体现。

TensorFlow py_function 嵌套输出类型？

TensorFlow py_function nested output type?

python

tensorflow

tensorflow-datasets