tf.assign和tf.while_loop不为人知的一面

tensorflow只运行相关节点

思考下面的代码，首先创建一个初始值为0变量aa，对变量aa赋值为1,然后把变量aa和常数2相加得到结果cc的值是多少？

import tensorflow as tf

aa = tf.Variable(0)
bb = tf.assign(aa, 1)
cc = aa + 2
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    print(sess.run(cc))

给人的第一反应，cc的值的应该为3，但结果似乎有它自己的想法。

运行结果如下：

2

tensorflow是声明式编程语言和传统python命令式编程有着较大的区别，tensorflow运行机制的本质是构建一幅有向无环图，然后往里面填充数据，运行程序。

可以把上例的图保存下来，然后用Tensorboard打开方便理解。

import tensorflow as tf

aa = tf.Variable(0)
bb = tf.assign(aa, 1)
cc = aa + 2
writer = tf.summary.FileWriter("./log",tf.get_default_graph())
writer.close()

通过命令行运行Tensorboard，即可在浏览器中输入地址http://127.0.0.1:6006查看图。(–host为指定显示地址)

tensorboard --logdir=./log --host=127.0.0.1

可以看出只运行了右边的分支。

得出结论：sess.run只会运行相关的节点，sess.run(cc)的时候，赋值操作并不会发生。

tf.assign的返回值

tf.assign是tensorflow中对变量赋值的节点，和普通的计算节点一样，tf.assign也有返回值，但不一样的是tf.assign的返回值是数据类型为dtype_ref的tensor，简单的说dtype_ref代表着数据类型为dtype变量的地址，当利用这个tensor的返回值进行其他计算时，会从变量的地址中取值参与计算。

import tensorflow as tf

aa = tf.Variable(0)
bb = tf.assign(aa, 1)
dd = tf.constant(0)

print(aa,bb,dd,sep='\n')

运行结果：

<tf.Variable 'Variable:0' shape=() dtype=int32_ref>
Tensor("Assign:0", shape=(), dtype=int32_ref)
Tensor("Const:0", shape=(), dtype=int32)

看下面这个例子。

import tensorflow as tf

aa = tf.Variable(0)
bb = tf.assign(aa, 1)
cc = tf.assign(bb, 2)
dd = cc + 3
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    for _ in range(10):
        print(sess.run((aa,bb,cc,dd)))

运行结果：

(2, 2, 2, 5)
(2, 2, 2, 5)
(2, 2, 2, 5)
(2, 2, 2, 5)
(2, 2, 2, 5)
(2, 2, 2, 5)
(2, 2, 2, 5)
(2, 2, 2, 5)
(2, 2, 2, 5)
(2, 2, 2, 5)

• 为什么aa,bb,cc的值相同？

(我的理解)tensorflow每次sess.run都会运行完相关节点的运算，最后从变量或者tensor里面取值得到最终结果，因为aa，bb，cc都只是保存着变量aa的地址，所以最终都是通过从变量aa的地址中取值，所以得到的结果是相同的。

• 利用tf.assign返回的tensor进行计算会产生依赖关系吗？

第一个tf.assign的返回值是int32_ref，里面保存的是变量aa的地址，把bb传入tf.assign进行赋值，这样做会产生依赖关系（先把变量aa赋值为1，第二个tf.assign才会对aa重新赋值2），从dd可以看出先把变量aa赋值为2，然后从aa的地址中取值出来参与计算。

tensorflow是并行计算的有向无环图

import tensorflow as tf
import time

aa = tf.Variable(0)
bb = tf.assign(aa, 1)
cc = tf.assign(aa, 2)
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    for _ in range(10):
        time.sleep(0.1)
        print(sess.run((aa,bb,cc)))

运行结果：

(1, 1, 1)
(1, 1, 1)
(1, 1, 1)
(1, 1, 1)
(1, 1, 1)
(1, 1, 1)
(2, 2, 2)
(1, 1, 1)
(1, 1, 1)
(2, 2, 2)

从结果可以看出，对变量赋值的先后顺序并不确定，这是因为由于tensorflow内部多线程运行的机制，导致bb和cc的先后顺序并不确定，所以要学会利用依赖关系建图和充分利用tensorflow的多线程机制加速运算。

tf.while_loop的运行机制

我的理解：
帮助文档中有一句话很关键（Note that while_loop calls cond and body exactly once ），指出cond和body函数只建了一次图。
tf.while_loop分别根据cond，和body函数建立了两个子图，然后把这两个子图和主图缝合起来。

import tensorflow as tf
import time

t_var = tf.Variable(0.0)
def cond(i, _):
    return tf.less(i, 3)
def body(i, _):
    global t_var
    flow = tf.assign(t_var, tf.cast(i, tf.float32))
    return tf.add(i, 1), flow
i, re = tf.while_loop(cond, body, [0, t_var])
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    for _ in range(10):
        time.sleep(0.1)
        print(sess.run([i, re]))

运行结果：

[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 1.0]
[3, 1.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]

根据前文的分析可知，re中储存的是变量t_var的地址，sess.run(re)即打印t_var的最终值，但是通过运行结果发现t_var的最终值并不确定。
重点来了！！！
第一次进行cond判断后返回True，进入body子图，此时运行机制类似与sess.run(tf.add(i, 1), flow)（单独建立一个body图，然后sess.run)。因为tensorflow的多线程机制，所以tf.add(i,1)可能在flow之前得到结果，由于tf.while_loop进行下一次循环只依赖tf.add(i,1)，所以可能不等上一个循环中flow节点的运行，下一个循环中的flow节点就运行了，这就是导致t_var的最终值不确定罪魁祸首。
虽然tf.while_loop只建了一次图，但它的表现类似于把多个body图拼接起来，如下图。

把线程参数parallel_iterations设置为1后（关闭多线程）就不会出现出现此种情况。

把
i, re = tf.while_loop(cond, body, [1, t_var])
替换为
i, re = tf.while_loop(cond, body, [1, t_var],parallel_iterations=1)

运行结果为：

[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]
[3, 2.0]

明确使用依赖关系同样可以避免赋值的先后顺序不明确。

import tensorflow as tf
import time

t_var = tf.Variable(0.0)
def cond(i, _):
    return tf.less(i, 6)
def body(i, flow):
    global t_var
    flow = 0. / (flow + 1.) + tf.cast(i, tf.float32)
    flow = tf.assign(t_var, flow)
    return tf.add(i, 1), flow
i, re = tf.while_loop(cond, body, [1, t_var])
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    for _ in range(10):
        time.sleep(0.1)
        print(sess.run([i, re]))

运行结果：

[6, 5.0]
[6, 5.0]
[6, 5.0]
[6, 5.0]
[6, 5.0]
[6, 5.0]
[6, 5.0]
[6, 5.0]
[6, 5.0]
[6, 5.0]

这样下一次赋值操作就会依赖上一次的赋值操作，就不会发生赋值顺序的不明确。

参考文章：
tensorflow文档
ziliwangmoe的博客