tf.while_loop

Posted 2021-01-26 wzdly

tf.while_loop(cond,
　　　　　　　　body,
　　　　　　　　loop_vars,
　　　　　　　　shape_invariants=None, 
　　　　　　　　parallel_iterations=10,
　　　　　　　　back_prop=True,
    　　　　　 swap_memory=False,
　　　　　　　　name=None,　　　　
　　　　　　　　maximum_iterations=None,
　　　　　　　　return_same_structure=False)

tf.while_loop 可以这样理解：

loop = []
while cond(loop):
    loop = body(loop)

即loop参数先传入cond 判断条件是否成立，成立之后，把 loop参数传入body 执行操作， 然后返回 操作后的 loop 参数，即loop参数已被更新，再把更新后的参数传入cond, 依次循环，直到不满足条件。

我们来看这样一个场景如何在 tensorflow中实现：

i= 0
n =10
while(i < n):
    i = i +1

首先这个要有个判断条件的语句，即：

i < n

满足条件就执行循环体里的操作，这个判断条件在tensorflow里，要写个函数来代替，即：

def cond(i, n):
    return i < n

之后是循环体里的操作，也要一个函数来代替，即：

def body(i, n):
    i = i + 1
    return i, n

请注意body函数里虽然没有与参数 n 有关的操作，但是必须要传入参数 n, 因为正如前面所说，要构成循环，参数在body函数更新后还要返回给cond函数，判断是否满足条件，如果不传入参数 n 下次，就没法判断了，合起来总得代码为：

import tensorflow as tf 

i = tf.constant(0)

c = lambda i: tf.less(i, 10)

b = lambda i: tf.add(i, 1)

r = tf.while_loop(c, b, [i])

with tf.Session() as sess:
    tf.global_variables_initializer().run()    
    print(sess.run(r))