TensorFlow py_function 嵌套输出类型？答案

【问题标题】：TensorFlow py_function nested output type?TensorFlow py_function 嵌套输出类型？
【发布时间】：2020-02-18 22:27:35
【问题描述】：

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

作为特定情况，我希望py_function 的返回类型为((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)) 类型，这将导致更干净的设置。

【问题讨论】：

标签： python tensorflow tensorflow-datasets

【解决方案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 中得到了证明。

【讨论】：