Tensorflow 2.0 中的 tf.function 和 tf.while 循环

Question

我正在尝试使用tf.while_loop 并行化循环。正如here 所建议的那样，parallel_iterations 参数在急切模式中没有任何区别。所以我尝试用tf.function 包裹tf.while_loop。但是，添加装饰器后，迭代变量的行为发生了变化。

例如，这段代码有效。

result = np.zeros(10)
iteration = tf.constant(0)
c = lambda i: tf.less(i, 10)
def print_fun(iteration):
    result[iteration] = iteration
    iteration += 1
    return (iteration,)
tf.while_loop(c, print_fun, [iteration])

如果我添加装饰器，就会出现错误。

result = np.zeros(10)
iteration = tf.constant(0)
c = lambda i: tf.less(i, 10)
def print_fun(iteration):
    result[iteration] = iteration
    iteration += 1
    return (iteration,)

@tf.function
def run_graph():
    iteration = tf.constant(0)
    tf.while_loop(c, print_fun, [iteration])

run_graph()

在调试过程中，我发现变量iteration 从张量变为占位符。这是为什么？我应该如何修改代码以消除该错误？

谢谢。

Answer 1

您的第一个 sn-p（没有 @tf.function 的那个）中的代码利用 TensorFlow 2 的急切执行来直接操作一个 numpy 数组（即您的外部 iteration 对象）。使用@tf.function，这不起作用，因为@tf.function 试图将您的代码编译成tf.Graph，它不能直接对numpy 数组进行操作（它只能处理tensorflow 张量）。要解决此问题，请使用 tf.Variable 并继续为其切片分配值。

使用@tf.function，您实际上可以使用更简单的代码来实现，利用@tf.function 的自动Python 到图形转换功能（称为AutoGraph）。您只需编写一个普通的 Python while 循环（使用 tf.less() 代替 < 运算符），while 循环将由 AutoGraph 编译成 tf.while_loop 在引擎盖下。

代码看起来像：

result = tf.Variable(np.zeros([10], dtype=np.int32))

@tf.function
def run_graph():
  i = tf.constant(0, dtype=tf.int32)
  while tf.less(i, 10):
    result[i].assign(i)  # Performance may require tuning here.
    i += 1

run_graph()
print(result.read_value())