Thursday, October 18, 2018

[TensorFlow Grappler] How to do the topological sorting in TensorFlow Grappler?

If you try to implement some optimizers in TensorFlow Grappler, you must have to know how to deal with the directed computation graph. One of the most important tools/knowledges is topological sorting.
The definition from Wiki: Topological sorting
https://en.wikipedia.org/wiki/Topological_sorting
"In the field of computer science, a topological sort or topological ordering of a directed graph is a linear ordering of its vertices such that for every directed edge uv from vertex u to vertex v, u comes before v in the ordering."



In TensorFlow, there are topological sort related functions already in the following link and we can take advantage of them.
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/grappler/utils/topological_sort.cc
First of all, we have this directed computation graph and can see the node: "fc1/fc1/MatMul"


Here is an example to dump the order of nodes in the graph using topological sorting
static std::unordered_map<const NodeDef*, int> GetTopoOrdering(GrapplerItem* item) {
  std::unordered_map<const NodeDef*, int> topo_order;
  ComputeTopologicalOrder(item->graph, &topo_order, nullptr);
  for ( auto& n : topo_order){
    const string& node_name = n.first->name();
    const int order = n.second;
    VLOG(1) << "...[DEBUG2] Node " << node_name << " at TopoOrdering order " << order;
  }
  return topo_order;
}

Result:
...[DEBUG2] Node train/Adam at TopoOrdering order 130
...[DEBUG2] Node train/Adam/Assign at TopoOrdering order 128
...[DEBUG2] Node train/Adam/mul at TopoOrdering order 126
...[DEBUG2] Node train/Adam/update_conv1/bias/ApplyAdam at TopoOrdering order 124
...[DEBUG2] Node train/gradients/conv1/BiasAdd_grad/BiasAddGrad at TopoOrdering order 122
...[DEBUG2] Node train/Adam/update_conv2/kernel/ApplyAdam at TopoOrdering order 121
...[DEBUG2] Node train/gradients/conv1/Relu_grad/ReluGrad at TopoOrdering order 120
...[DEBUG2] Node train/Adam/update_conv1/kernel/ApplyAdam at TopoOrdering order 125
...[DEBUG2] Node train/gradients/conv2/Conv2D_grad/Conv2DBackpropFilter at TopoOrdering order 119
...[DEBUG2] Node train/gradients/conv2/Relu_grad/ReluGrad at TopoOrdering order 115
...[DEBUG2] Node train/gradients/pool3/dropout/cond/dropout/div_grad/tuple/control_dependency at TopoOrdering order 109
...[DEBUG2] Node train/gradients/pool3/dropout/cond/dropout/div_grad/RealDiv at TopoOrdering order 108
...[DEBUG2] Node train/gradients/conv1/Conv2D_grad/Conv2DBackpropFilter at TopoOrdering order 123
...[DEBUG2] Node train/gradients/pool3/dropout/cond/Identity/Switch_grad/cond_grad at TopoOrdering order 106
...[DEBUG2] Node train/Adam/update_fc1/kernel/ApplyAdam at TopoOrdering order 105
...[DEBUG2] Node train/gradients/pool3/MaxPool_grad/MaxPoolGrad-2-TransposeNHWCToNCHW-LayoutOptimizer at TopoOrdering order 113
...[DEBUG2] Node train/gradients/pool3/dropout/cond/Merge_grad/cond_grad at TopoOrdering order 104
...[DEBUG2] Node train/Adam/update_output/kernel/ApplyAdam at TopoOrdering order 99
...[DEBUG2] Node train/gradients/fc1/fc1/Relu_grad/ReluGrad at TopoOrdering order 98
...[DEBUG2] Node train/gradients/output/output/MatMul_grad/MatMul at TopoOrdering order 95
...[DEBUG2] Node train/gradients/conv2/BiasAdd_grad/BiasAddGrad at TopoOrdering order 116
...[DEBUG2] Node train/gradients/output/output/BiasAdd_grad/BiasAddGrad at TopoOrdering order 94
...[DEBUG2] Node output/output/BiasAdd at TopoOrdering order 90
...[DEBUG2] Node output/output/MatMul at TopoOrdering order 89
...[DEBUG2] Node fc1/fc1/BiasAdd at TopoOrdering order 87
...[DEBUG2] Node fc1/fc1/Relu at TopoOrdering order 88
...[DEBUG2] Node pool3/dropout/cond/Merge at TopoOrdering order 85
...[DEBUG2] Node pool3/dropout/cond/dropout/mul at TopoOrdering order 82
...[DEBUG2] Node train/gradients/fc1/fc1/MatMul_grad/MatMul at TopoOrdering order 101
...[DEBUG2] Node train/gradients/zeros/Const at TopoOrdering order 80
...[DEBUG2] Node train/gradients/Shape_1 at TopoOrdering order 79
...[DEBUG2] Node pool3/dropout/cond/dropout/div at TopoOrdering order 78
...[DEBUG2] Node ConstantFoldingCtrl/pool3/dropout/cond/dropout/div/Switch_1 at TopoOrdering order 77
...[DEBUG2] Node train/gradients/Switch at TopoOrdering order 76
...[DEBUG2] Node pool3/dropout/cond/dropout/div/Switch at TopoOrdering order 75
...[DEBUG2] Node pool3/MaxPool-1-0-TransposeNCHWToNHWC-LayoutOptimizer at TopoOrdering order 73
...[DEBUG2] Node pool3/dropout/cond/dropout/Floor at TopoOrdering order 72
...[DEBUG2] Node pool3/MaxPool at TopoOrdering order 71
...[DEBUG2] Node ConstantFolding/train/gradients/conv2/Conv2D_grad/ShapeN-matshapes-1 at TopoOrdering order 67
...[DEBUG2] Node conv2/BiasAdd at TopoOrdering order 66
...[DEBUG2] Node train/gradients/pool3/dropout/cond/dropout/mul_grad/Mul at TopoOrdering order 107
...[DEBUG2] Node train/gradients/conv2/Conv2D_grad/ShapeN at TopoOrdering order 64
...[DEBUG2] Node pool3/dropout/cond/dropout/random_uniform/RandomUniform at TopoOrdering order 62
...[DEBUG2] Node pool3/dropout/cond/dropout/Shape at TopoOrdering order 59
...[DEBUG2] Node train/gradients/conv2/Conv2D_grad/Conv2DBackpropInput at TopoOrdering order 118
...[DEBUG2] Node pool3/dropout/cond/dropout/keep_prob at TopoOrdering order 58
...[DEBUG2] Node conv1/BiasAdd at TopoOrdering order 57
...[DEBUG2] Node ConstantFolding/train/gradients/conv1/Conv2D_grad/ShapeN-matshapes-1 at TopoOrdering order 54
...[DEBUG2] Node train/gradients/conv1/Conv2D_grad/Conv2DBackpropFilter-0-TransposeNHWCToNCHW-LayoutOptimizer at TopoOrdering order 52
...[DEBUG2] Node conv1/Conv2D-0-TransposeNHWCToNCHW-LayoutOptimizer at TopoOrdering order 51
...[DEBUG2] Node train/Adam/mul_1 at TopoOrdering order 127
...[DEBUG2] Node train/beta2_power/read at TopoOrdering order 50
...[DEBUG2] Node train/gradients/zeros_1 at TopoOrdering order 84
...[DEBUG2] Node train/beta1_power/read at TopoOrdering order 49
...[DEBUG2] Node train/Adam/update_fc1/bias/ApplyAdam at TopoOrdering order 103
...[DEBUG2] Node train/gradients/output/output/MatMul_grad/MatMul_1 at TopoOrdering order 96
...[DEBUG2] Node output/bias/read at TopoOrdering order 48
...[DEBUG2] Node output/kernel/read at TopoOrdering order 47
...[DEBUG2] Node fc1/bias/read at TopoOrdering order 46
...[DEBUG2] Node fc1/kernel/read at TopoOrdering order 45
...[DEBUG2] Node conv2/kernel/read at TopoOrdering order 43
...[DEBUG2] Node conv1/bias/read at TopoOrdering order 42
...[DEBUG2] Node conv1/kernel/read at TopoOrdering order 41
...[DEBUG2] Node train/gradients/fc1/fc1/MatMul_grad/MatMul_1 at TopoOrdering order 102
...[DEBUG2] Node inputs/training at TopoOrdering order 40
...[DEBUG2] Node inputs/Reshape at TopoOrdering order 39
...[DEBUG2] Node PermConstNCHWToNHWC-LayoutOptimizer at TopoOrdering order 38
...[DEBUG2] Node pool3/dropout/cond/dropout/random_uniform at TopoOrdering order 68
...[DEBUG2] Node PermConstNHWCToNCHW-LayoutOptimizer at TopoOrdering order 37
...[DEBUG2] Node train/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits at TopoOrdering order 91
...[DEBUG2] Node train/Adam/epsilon at TopoOrdering order 36
...[DEBUG2] Node conv2/Conv2D at TopoOrdering order 63
...[DEBUG2] Node train/Adam/beta2 at TopoOrdering order 35
...[DEBUG2] Node train/Adam/update_output/bias/ApplyAdam at TopoOrdering order 97
...[DEBUG2] Node train/Adam/beta1 at TopoOrdering order 34
...[DEBUG2] Node train/gradients/train/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits_grad/PreventGradient at TopoOrdering order 92
...[DEBUG2] Node pool3/Reshape at TopoOrdering order 74
...[DEBUG2] Node conv2/Relu at TopoOrdering order 69
...[DEBUG2] Node output/bias/Adam_1 at TopoOrdering order 32
...[DEBUG2] Node train/gradients/zeros at TopoOrdering order 83
...[DEBUG2] Node output/bias/Adam at TopoOrdering order 31
...[DEBUG2] Node output/kernel/Adam_1 at TopoOrdering order 30
...[DEBUG2] Node output/kernel/Adam at TopoOrdering order 29
...[DEBUG2] Node fc1/fc1/MatMul at TopoOrdering order 86
...[DEBUG2] Node fc1/bias/Adam_1 at TopoOrdering order 28
...[DEBUG2] Node train/gradients/pool3/MaxPool_grad/MaxPoolGrad at TopoOrdering order 114
...[DEBUG2] Node conv1/Relu at TopoOrdering order 61
...[DEBUG2] Node fc1/bias/Adam at TopoOrdering order 27
...[DEBUG2] Node train/gradients/pool3/Reshape_grad/Reshape at TopoOrdering order 112
...[DEBUG2] Node fc1/kernel/Adam_1 at TopoOrdering order 26
...[DEBUG2] Node train/gradients/pool3/dropout/cond/dropout/div/Switch_grad/cond_grad at TopoOrdering order 110
...[DEBUG2] Node pool3/dropout/cond/dropout/add at TopoOrdering order 70
...[DEBUG2] Node fc1/kernel/Adam at TopoOrdering order 25
...[DEBUG2] Node conv2/bias/Adam_1 at TopoOrdering order 24
...[DEBUG2] Node train/gradients/train/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits_grad/mul at TopoOrdering order 93
...[DEBUG2] Node train/gradients/Shape_2 at TopoOrdering order 81
...[DEBUG2] Node conv2/bias/Adam at TopoOrdering order 23
...[DEBUG2] Node train/Adam/update_conv2/bias/ApplyAdam at TopoOrdering order 117
...[DEBUG2] Node train/gradients/AddN at TopoOrdering order 111
...[DEBUG2] Node conv2/kernel/Adam_1 at TopoOrdering order 22
...[DEBUG2] Node conv2/kernel/Adam at TopoOrdering order 21
...[DEBUG2] Node pool3/dropout/cond/dropout/random_uniform/mul at TopoOrdering order 65
...[DEBUG2] Node conv1/bias/Adam_1 at TopoOrdering order 20
...[DEBUG2] Node conv1/kernel/Adam_1 at TopoOrdering order 18
...[DEBUG2] Node train/Adam/Assign_1 at TopoOrdering order 129
...[DEBUG2] Node conv1/kernel/Adam at TopoOrdering order 17
...[DEBUG2] Node train/beta2_power at TopoOrdering order 16
...[DEBUG2] Node train/beta1_power at TopoOrdering order 15
...[DEBUG2] Node train/gradients/pool3/Reshape_grad/Shape at TopoOrdering order 14
...[DEBUG2] Node train/gradients/train/SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits_grad/ExpandDims at TopoOrdering order 13
...[DEBUG2] Node pool3/dropout/cond/switch_t at TopoOrdering order 56
...[DEBUG2] Node conv1/bias/Adam at TopoOrdering order 19
...[DEBUG2] Node output/bias at TopoOrdering order 12
...[DEBUG2] Node fc1/bias at TopoOrdering order 10
...[DEBUG2] Node train/Adam/learning_rate at TopoOrdering order 33
...[DEBUG2] Node fc1/kernel at TopoOrdering order 9
...[DEBUG2] Node pool3/Reshape/shape at TopoOrdering order 8
...[DEBUG2] Node conv2/bias at TopoOrdering order 7
...[DEBUG2] Node pool3/dropout/cond/Switch at TopoOrdering order 53
...[DEBUG2] Node conv2/kernel at TopoOrdering order 6
...[DEBUG2] Node conv1/bias at TopoOrdering order 5
...[DEBUG2] Node conv1/kernel at TopoOrdering order 4
...[DEBUG2] Node conv1/Conv2D at TopoOrdering order 55
...[DEBUG2] Node inputs/training/input at TopoOrdering order 3
...[DEBUG2] Node output/kernel at TopoOrdering order 11
...[DEBUG2] Node inputs/y at TopoOrdering order 2
...[DEBUG2] Node train/gradients/fc1/fc1/BiasAdd_grad/BiasAddGrad at TopoOrdering order 100
...[DEBUG2] Node ConstantFolding/pool3/dropout/cond/dropout/div_recip at TopoOrdering order 60
...[DEBUG2] Node inputs/Reshape/shape at TopoOrdering order 1
...[DEBUG2] Node conv2/bias/read at TopoOrdering order 44
...[DEBUG2] Node inputs/X at TopoOrdering order 0


So, we can see the sequence of the following nodes in the directed computation graph:
fc1/fc1/MatMul ==> fc1/fc1/BiasAdd ==> fc1/fc1/Relu
And the topological sorting method gives us the same sequence

...[DEBUG2] Node fc1/fc1/BiasAdd at TopoOrdering order 87
...[DEBUG2] Node fc1/fc1/Relu at TopoOrdering order 88
...[DEBUG2] Node fc1/fc1/MatMul at TopoOrdering order 86

No comments: