不允许迭代 tf.Tensor：此函数中禁用 AutoGraph

Question

我正在使用 tensorflow 2.1 和 python 3.7

以下代码片段用于构建张量流图。当作为独立的 python 脚本执行时，代码运行没有错误。（可能 tensorflow 正在以 Eager 模式运行？我不确定。）

import tensorflow as tf
patches = tf.random.uniform(shape=(1, 10, 50, 300), dtype=tf.dtypes.float32)
s = tf.shape(patches)
patches = [patches[0][x][y] - tf.reduce_mean(patches[0][x][y]) for y in tf.range(s[2]) for x in tf.range(s[1])]

但是，当这是张量流图的一部分时，代码会失败。我收到以下错误：tensorflow。

python.framework.errors_impl.OperatorNotAllowedInGraphError：tf.Tensor不允许迭代：此函数中禁用了 AutoGraph。尝试直接用@tf.function 装饰它。

我还将装饰器添加@tf.function到包装上述代码行的方法中。它没有帮助。我不确定我是否完全理解用 装饰的含义@tf.function。我还检查了这可能是在张量流图中使用 python 列表理解的问题。我不确定如何使用tf.map_fn或tf.while_loop来处理我的情况，因为我有嵌套循环。

提前致谢！

Answer 1

亲笔签名尚不支持列表推导式。提出的错误也需要改进。堆积在https://github.com/tensorflow/tensorflow/issues/32546应该有助于尽快解决它。

在支持理解之前，您必须使用 map_fn，在这种情况下，它看起来像这样：

def outer_comp(x):
  def inner_comp(y):
    return patches[0][x][y] - tf.reduce_mean(patches[0][x][y])
  return tf.map_fn(inner_comp, tf.range(s[2]), dtype=tf.float32)
patches = tf.map_fn(outer_comp, tf.range(s[1]), dtype=tf.float32)

也就是说，我相信你可以直接使用 reduce_mean ：

patches = patches - tf.expand_dims(tf.reduce_mean(patches, axis=3), -1)