Incubator-mxnet: How do I interpret the JSON compute graph?

[aidan-plenert-macdonald]

I am trying to interpret the symbolic compute graph that MxNet is dumping out.

When I run the code,

import mxnet as mx
from mxnet.gluon import nn
from mxnet import nd

net = nn.HybridSequential()
with net.name_scope():
    net.add(nn.Dense(256, activation="relu"))
    net.add(nn.Dense(128, activation="relu"))
    net.add(nn.Dense(2))
net.collect_params().initialize()

x = mx.sym.var('input')
y = net(x)

print(y.tojson())

I get a large JSON structure with things like,

{
      "op": "FullyConnected", 
      "name": "hybridsequential1_dense0_fwd", 
      "attr": {"num_hidden": "256"}, 
      "inputs": [[0, 0, 0], [1, 0, 0], [2, 0, 0]]
}

full json

{
  "nodes": [
    {
      "op": "null", 
      "name": "input", 
      "inputs": []
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense0_weight", 
      "attr": {
        "__dtype__": "0", 
        "__lr_mult__": "1.0", 
        "__shape__": "(256, 0)", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense0_bias", 
      "attr": {
        "__dtype__": "0", 
        "__init__": "zeros", 
        "__lr_mult__": "1.0", 
        "__shape__": "(256,)", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "FullyConnected", 
      "name": "hybridsequential1_dense0_fwd", 
      "attr": {"num_hidden": "256"}, 
      "inputs": [[0, 0, 0], [1, 0, 0], [2, 0, 0]]
    }, 
    {
      "op": "Activation", 
      "name": "hybridsequential1_dense0_relu_fwd", 
      "attr": {"act_type": "relu"}, 
      "inputs": [[3, 0, 0]]
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense1_weight", 
      "attr": {
        "__dtype__": "0", 
        "__lr_mult__": "1.0", 
        "__shape__": "(128, 0)", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense1_bias", 
      "attr": {
        "__dtype__": "0", 
        "__init__": "zeros", 
        "__lr_mult__": "1.0", 
        "__shape__": "(128,)", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "FullyConnected", 
      "name": "hybridsequential1_dense1_fwd", 
      "attr": {"num_hidden": "128"}, 
      "inputs": [[4, 0, 0], [5, 0, 0], [6, 0, 0]]
    }, 
    {
      "op": "Activation", 
      "name": "hybridsequential1_dense1_relu_fwd", 
      "attr": {"act_type": "relu"}, 
      "inputs": [[7, 0, 0]]
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense2_weight", 
      "attr": {
        "__dtype__": "0", 
        "__lr_mult__": "1.0", 
        "__shape__": "(2, 0)", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "null", 
      "name": "hybridsequential1_dense2_bias", 
      "attr": {
        "__dtype__": "0", 
        "__init__": "zeros", 
        "__lr_mult__": "1.0", 
        "__shape__": "(2,)", 
        "__wd_mult__": "1.0"
      }, 
      "inputs": []
    }, 
    {
      "op": "FullyConnected", 
      "name": "hybridsequential1_dense2_fwd", 
      "attr": {"num_hidden": "2"}, 
      "inputs": [[8, 0, 0], [9, 0, 0], [10, 0, 0]]
    }
  ], 
  "arg_nodes": [0, 1, 2, 5, 6, 9, 10], 
  "node_row_ptr": [
    0, 
    1, 
    2, 
    3, 
    4, 
    5, 
    6, 
    7, 
    8, 
    9, 
    10, 
    11, 
    12
  ], 
  "heads": [[11, 0, 0]], 
  "attrs": {"mxnet_version": ["int", 1100]}
}

How do I interpret the input? How do I know which nodes connect where?

[sergeykolychev]

@aidan-plenert-macdonald AFAIK this json is not intended to be easily readable by human and supposed to be used by mx.symbol.load function to recreate the graph.
That being said your net is pretty simple and that makes the json somewhat readable.
Go from the top and assign to each node an index starting from 0 as they appear from the top.
First node is input (index 0), second and third are weight and bias tensors for the first dense op (or in reality FullyConnected operation) which is fourth in the line (idx 3); this node has inputs node that tells you that it accepts three nodes as inputs 0, 1, 2 - input, weight, bias; Next in line (idx 4) is operation Activation (you created dense with relu activation) and it has input node 3 which is output from the FullyConnected op. This way you can follow the logic to the bottom. But the json is not intended for this usage. Use vizualization module of mxnet to create a png of the symbol structure.

[aidan-plenert-macdonald]

@sergeykolychev Thanks. I wanted to look at the JSON so I can create a Jupyter widget that displays the network with sliders and buttons to configure the network. Kind of like the PNG.

So why are the inputs 2D. What are the zeros after the node indexes? eg.

"inputs": [[4, 0, 0], [5, 0, 0], [6, 0, 0]]

[sergeykolychev]

@aidan-plenert-macdonald that's beyond my knowledge on the matter, sorry.

[zhreshold]

@aidan-plenert-macdonald
Here's what I wrote for parsing nnvm JSON file. http://nnvm.tvmlang.org/json_spec.html
MXNet's JSON follow a similar routine, though some of the operators may vary and and 'attrs' is 'attr' in mxnet.
More specifically to your problem, inputs are list of entries, which is [nose_id, index, version]
Hope this helps.

[Tveek]

@zhreshold , can you explain index and version？

[ChiaraXian]

@zhreshold Could you explain inputs of mxnet .json nodes in detail? Especially index and version. I am now trying to read mxnet .json and understand the node structure but I haven't found any document about this. Your answer will be really helpful for me. Thanks!

[aidan-plenert-macdonald]

@ChiaraXian Maybe look at the code I wrote here will help. It's been a while, but I might be able to help if you have questions

https://github.com/aidan-plenert-macdonald/mxflows/blob/master/mxflows/flow_diagram.py#L45

[ChiaraXian]

@aidan-plenert-macdonald Thank you for your code. But I still wonder the index and version in the list of inputs. [node_id, index, version]. The first item means the node_id of input, how about the other two items?

[aidan-plenert-macdonald]

Found the documentation here https://docs.tvm.ai/dev/nnvm_json_spec.html#nodes. It's not super helpful.

You can trace down the code a bit,

https://github.com/dmlc/nnvm/blob/master/python/nnvm/graph.py#L26 ==>
https://github.com/dmlc/nnvm/blob/master/python/nnvm/graph.py#L270 ==>
...
https://github.com/dmlc/nnvm/blob/ef0ab9b09dbf1318851be311d3752de6c9bd4881/src/core/graph.cc#L45 ==>
https://github.com/dmlc/nnvm/blob/ef0ab9b09dbf1318851be311d3752de6c9bd4881/include/nnvm/node.h#L31

struct NodeEntry {
  /*! \brief the source node of this data */
  NodePtr node;
  /*! \brief index of output from the source. */
  uint32_t index;
  /*!
   * \brief version of input Variable.
   *  This field can only be nonzero when this->node is a Variable node.
   *  version is increased by one each time a Variable get composed to a mutation Op.
   *  This information can be helpful to decide order of operations when sequence of mutation happens.
   */
  uint32_t version;
};

Not really sure what that means, but you could probably figure that out with some more digging

[ChiaraXian]

@aidan-plenert-macdonald Thanks a lot!!! Thank you for your help!!!