### Summary
In today’s TVM TensorFlow frontend, there is only limited support for control
flow, which resulting in difficult in covering TensorFlow object detection
models. In this RFC, we will discuss how to improve current TVM TF frontend to
fully support TensorFlow control flow.
### Solution
#### Visit TensorFlow Graph Def
Currently TVM parses TensorFlow Graph Def node by node in topological order.
However, topological order might not stand any more if control flow is
introduced into graph def. We do the following to alter the visiting order:
1. Visit all nodes in original graph def order and fetch all control flow nodes.
2. For all control flows nodes, modify the order so that for each while_loop,
all Exit nodes sit at the end.
3. Visit all control flow nodes in modified order and generate corresponding
Relay Functions. For each control flow node, we backtrack all its ancestor
nodes until input nodes.
4. Visit and parse other unvisited nodes.
#### Parse Control Flow nodes
In TVM, while loop is represented as a recursive function. To convert a
TensorFlow while_loop block into a Relay function, the following assumption
must stand:
1. All `Exit` nodes under the same while_loop must have the same node name
prefix as this while loop block, which is got with `node.name.rsplit('/',
1)[0]`. In TensorFlow 1.x, user can specify the name of while_loop, but `Exit`
node prefix is automated generated with ancestor nodes’ names.
2. `Merge`, `Switch`, `NextIteration` and `Exit` nodes with the same postfix
number are belong to the same iteration body. For example, `Merge_1`,
`Switch_1`, `NextIteration_1` and `Exit_1` should refer to the same iteration
body. Since these four nodes are generated when user creating while_loop, the
postfix number can’t be arbitrarily set.
We convert a while loop block to a Relay function with the following method:
1. When we get a `Merge` node inside a while_loop block: if it is a direct
child of a `while_loop` block, we create a Loop wrapper to start adding all
sub-nodes under this loop block. Otherwise, we create a complete branch op with
previously stored Switch condition.
2. When we get an `Enter` node: we simply convert its input node to Relay node.
3. When we get a `LoopCond` node: we first convert its input node to Relay
node. Then we set the converted node as our Loop wrapper condition.
4. When we get a `Switch` node: we first convert both its input and condition
to Relay nodes. If it is a direct child of a while_loop block, we add
converted input to loop variable list of the corresponding Loop wrapper.
Otherwise, we create a Branch wrapper, set its condition to be converted
condition node and add it to a global branch dictionary.
5. When we get an `Exit` node: we generate Relay function for the corresponding
while loop body.
### Tested models
With above changes(together with pending PR of dynamic NMS), we can convert and
compile TensorFlow object detection models. The following models from
TensorFlow object detection model zoo are tested:
ssd_mobilenet_v1_coco
ssd_mobilenet_v2_coco
ssd_resnet_50_fpn_coco
mask_rcnn_inception_v2_coco
mask_rcnn_resnet50_atrous_coco
faster_rcnn_resnet50_coco
###TODO
Investigate TensorFlow 2.0 control flow and what changes are required.
@tqchen @jroesch @srkreddy1238 @zhiics @yongwww @wweic
--
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/apache/incubator-tvm/issues/4969
