python - TF map_fn 或 while_loop 用于不同形状的张量列表

Question

我想处理不同形状的张量序列(列表)并输出另一个张量列表。考虑每个时间戳上具有不同隐藏状态大小的 RNN。类似的东西

输入:[tf.ones((1, 2, 2)), tf.ones((2, 2, 3)), tf.ones((3, 2, 1))]

输出:[tf.zeros((1, 2, 4)), tf.zeros((4, 2, 6)), tf.zeros((6, 2, 1))]

我无法将输入(或输出)堆叠到单个张量中，因为它们都有不同的形状，因此我无法使用 tf.map_fn 来完成该任务。目前，我使用 python for 循环，但它似乎不是最理想的。

我还能做些什么更好的事情吗？

Answer 1

您可以使用tf.while_loop重复执行任意 TensorFlow 操作，直到发生某些停止条件。停止条件本身被指定为一个操作。

请注意，应谨慎使用tf.while_loop，因为默认情况下它的迭代将并行运行。例如，如果循环体递增 tf.Variable，则必须使用 control dependencies确保迭代按顺序运行。

但是，您提到您有一个带有 Python 循环的有效实现。如果可能，使用 Python 进行循环通常是最有效的解决方案。当您在 Python 中构建循环时，您可以为循环中的每次迭代创建单独的操作。这让 TensorFlow 在图构建时决定如何为每个操作分配计算资源。例如，如果提前知道迭代次数，则更容易预测内存需求和并行化可能性。

因此，当构建图时未知停止条件时，最常使用 tf.while_loop 和 tf.map_fn。

如果迭代次数固定但非常大，您可能仍然想使用 tf.while_loop 而不是 Python 循环，因为每个操作的内存成本不菲。