tf.while_loop循环 实现

时间:2025-04-07 07:43:43

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 下次,就没法判断了。

合起来总得代码为

 i  = 0
 n =10 

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

def body(i, n):
    i = i + 1
    return i, n
i, n = tf.while_loop(cond, body, [i, n])

可以执行的代码

import tensorflow as tf 
i = tf.get_variable("ii", dtype=tf.int32, shape=[], initializer=tf.ones_initializer())
n = (10)

def cond(a, n):
    return  a< n
def body(a, n):
    a = a + 1
    return a, n

a, n = tf.while_loop(cond, body, [a, n])
with () as sess:
    tf.global_variables_initializer().run()
    res = ([a, n])
    print(res)