tf.cond and tf.whileloop
本篇博客梳理一下 tensorflow python client API 与tf.cond和 tf.whileloop相关的部分,对于这两个api 的详细解释,请看我的第一篇博客和第二篇博客,这篇博客讲一些api 的实现细节。
API 梳理
tf.cond
下图作为讲解tf.cond代码时的参考
以下是tf.cond的伪代码。源代码在controlf_flow_ops.py
context_t = CondContext(pred, branch=1)
res_t = context_t.Call(fn1)
# Build the graph for the false branch
context_f = CondContext(pred, branch=0)
res_f = context_f.Call(fn2)
# Add the Merge nodes for the outputs
merges = [Merge([f, t]) for (f, t) in zip(res_f, res_t)]
return merges
正如我在第二篇博客提到的那样。
tf.whileloop
下图作为tf.whileloop的参考:
以下是tf.whileloop的伪代码。源代码在controlf_flow_ops.py
while_context = WhileContext()
while_context.Enter()
# Add the Enter nodes for each loop variable.
enter_vars = [Enter(x, frame_name) for x in loop_vars]
# Add the Merge nodes. Note that input[1] will be updated later.
merge_vars = [Merge([x, x]) for x in enter_vars]
# Build the loop pred subgraph.
pred_result = pred(*merge_vars)
# Add the Switch nodes.
switch_vars = [Switch(x, pred_result) for x in merge_vars]
# Build the loop body subgraph.
body_result = body(*[x[1] for x in switch_vars])
# Add the NextIteration nodes.
next_vars = [NextIteration(x) for x in body_result]
# Form the cycles for the loop.
for m, v in zip(merge_vars, next_vars):
m.op._update_input(1, v)
# Add the Exit nodes.
exit_vars = [Exit(x[0]) for x in switch_vars]
while_context.Exit()
return exit_vars
实现细节
由上面的伪代码可知,tf.cond和tf.whileloop实现的时候都是建立一个上下文(context),由此抽象出一个基类ControlFlowContext,来处理cond和whileloop相关的地方。一个控制流上下文主要保存了控制流相关的上下文信息,当在一个上下文里新建一个operation ,operation初始化时会用到这些上下文。比如说,如果在一个whileloop里新建一个operation,而这个whileloop是一个前向传播whileloop的反向传播(详情请参考我的第二篇博客),而且这个operation 需要用到前向传播的迭代中计算出来的值。前先传播中值的生产和反向传播中对值的消费,是一种first in last out 的关系,所以要加入一个栈和相应的入栈和出栈操作,如下图所示。
此外,从一个控制流上下文输出的值也需要做处理、对控制流上下文里op控制依赖边也要做一些调整。这些都依赖控制上下文信息。
class ControlFlowContext(object)
这个基类抽象出对cond 和whileloop 通用的接口。
以下分别介绍cond 和whileloop 两个类。
由上面的两张图,我们可以知道,cond 和 whileloop 的一个重要的相同之处是,它们都要对在这个上下文中创建的op的输入做一些处理,这些处理由addop这个方法完成。下面我主要讲这个方法
这里面的逻辑是,一个op 属于某一个controflow context, controflow context是嵌套的,当一个op 初始化时,对该op 的输入值,调用op 所在ControlFlowContext的addvalues 方法。对于cond,是以传入的输入值为基础添加一个merge节点;对于whileloop,可能需要加入入栈和出栈的操作。
__init__ (python/framework/ops.py 里的operation初始化方法)
-
self._control_flow_post_processing
-
self._control_flow_context.AddOp(self)
-
class CondContext(ControlFlowContext)
class WhileContext(ControlFlowContext)