tf.while_loop

参考  tf.while_loop - 云+社区 - 腾讯云

当条件为真时，重复身体动作。

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

cond是一个返回布尔标量张量的可调用的张量。body是一个可调用的变量，返回一个(可能是嵌套的)元组、命名元组或一个与loop_vars具有相同特性(长度和结构)和类型的张量列表。loop_vars是一个(可能是嵌套的)元组、命名元组或张量列表，它同时传递给cond和body。cond和body都接受与loop_vars一样多的参数。除了常规张量或索引片之外，主体还可以接受和返回TensorArray对象。TensorArray对象的流将在循环之间和梯度计算期间适当地转发。注意while循环只调用cond和body一次(在调用while循环的内部调用，而在Session.run()期间根本不调用)。while loop使用一些额外的图形节点将cond和body调用期间创建的图形片段拼接在一起，创建一个图形流，该流重复body，直到cond返回false。为了保证正确性，tf.while循环()严格地对循环变量强制执行形状不变量。形状不变量是一个(可能是部分的)形状，它在循环的迭代过程中保持不变。如果循环变量的形状在迭代后被确定为比其形状不变量更一般或与之不相容，则会引发错误。例如，[11,None]的形状比[11,17]的形状更通用，而且[11,21]与[11,17]不兼容。默认情况下(如果参数shape_constant没有指定)，假定loop_vars中的每个张量的初始形状在每次迭代中都是相同的。shape_constant参数允许调用者为每个循环变量指定一个不太特定的形状变量，如果形状在迭代之间发生变化，则需要使用该变量。tf.Tensor。体函数中也可以使用set_shape函数来指示输出循环变量具有特定的形状。稀疏张量和转位切片的形状不变式特别处理如下:

a)如果一个循环变量是稀疏张量，那么形状不变量必须是张量形状([r])，其中r是由稀疏张量表示的稠密张量的秩。这意味着稀疏张量的三个张量的形状是([None]， [None, r]， [r])。注意:这里的形状不变量是SparseTensor.dense_shape属性的形状。它一定是向量的形状。

b)如果循环变量是索引切片，则形状不变量必须是索引切片的值张量的形状不变量。它表示索引切片的三个张量的形状为(shape， [shape[0]]， [shape.ndims])。

while循环实现了非严格的语义，允许多个迭代并行运行。并行迭代的最大数量可以由parallel_iteration控制，这让用户可以控制内存消耗和执行顺序。对于正确的程序，while循环应该为任何parallel_iteration > 0返回相同的结果。

对于训练，TensorFlow存储了在正向推理中产生的、在反向传播中需要的张量。这些张量是内存消耗的主要来源，在gpu上进行训练时经常导致OOM错误。当swap_memory标志为true时，我们将这些张量从GPU交换到CPU。例如，这允许我们用很长的序列和大量训练RNN模型。

参数：

• cond:表示循环终止条件的可调用的。
• body:表示循环体的可调用的。
• loop_vars:一个(可能是嵌套的)元组、命名元组或numpy数组、张量和TensorArray对象列表。
• shape_constant:循环变量的形状不变量。
• parallel_iteration:允许并行运行的迭代数。它必须是一个正整数。
• back_prop:是否为此while循环启用backprop。
• swap_memory:这个循环是否启用了GPU-CPU内存交换。
• maximum_iteration:可选的while循环运行的最大迭代次数。如果提供了cond输出，则使用附加条件来确保执行的迭代数不大于maximum_iteration。
• name:返回的张量的可选名称前缀。

返回值：

• 循环变量的输出张量。返回值具有与loop_vars相同的结构。

可能产生的异常：

• TypeError: if cond or body is not callable.
• ValueError: if loop_vars is empty.

例：

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])

嵌套和命名元组的例子：

import collections
Pair = collections.namedtuple('Pair', 'j, k')
ijk_0 = (tf.constant(0), Pair(tf.constant(1), tf.constant(2)))
c = lambda i, p: i < 10
b = lambda i, p: (i + 1, Pair((p.j + p.k), (p.j - p.k)))
ijk_final = tf.while_loop(c, b, ijk_0)

使用形状不变量的例子：

i0 = tf.constant(0)
m0 = tf.ones([2, 2])
c = lambda i, m: i < 10
b = lambda i, m: [i+1, tf.concat([m, m], axis=0)]
tf.while_loop(
    c, b, loop_vars=[i0, m0],
    shape_invariants=[i0.get_shape(), tf.TensorShape([None, 2])])

示例演示了不严格的语义:在下面的示例中,计数器的最终值不依赖于x,所以while_loop可以增加与x的更新并行的计数器,但是,因为一个循环迭代中的循环计数器取决于之前迭代的值,循环计数器本身不能并行地递增。因此,如果我们只是想要计数器的最终值(我们在行打印上打印(sess.run(i)),那么x将永远不会递增,但是计数器将在一个线程上更新。相反,如果我们想要输出的值(我们在行打印上打印(sess.run(out). shape),那么计数器可能会在自己的线程上递增,而x可以在一个单独的线程上并行地递增。在极端情况下,可以想象的是,在x增加到一个时间之前,线程会将计数器运行,直到完成。唯一不能发生的事情是,线程更新x永远不能超过计数器线程,因为线程递增x取决于计数器的值。

import tensorflow as tf
 
n = 10000
x = tf.constant(list(range(n)))
c = lambda i, x: i < n
b = lambda i, x: (tf.compat.v1.Print(i + 1, [i]), tf.compat.v1.Print(x + 1,
[i], "x:"))
i, out = tf.while_loop(c, b, (0, x))
with tf.compat.v1.Session() as sess:
    print(sess.run(i))  # prints [0] ... [9999]
 
    # The following line may increment the counter and x in parallel.
    # The counter thread may get ahead of the other thread, but not the
    # other way around. So you may see things like
    # [9996] x:[9987]
    # meaning that the counter thread is on iteration 9996,
    # while the other thread is on iteration 9987
    print(sess.run(out).shape)