Added edges to graph. Also added script to download mock models (#39)

* Add edges to graph. Add squeezenet to mock data.

* rm json

* rm pb files and add links to them

* Add edges to graph. Add squeezenet to mock data.

rm json

rm pb files and add links to them

* added download script

* typo

server/visualdl/graph.py

@@ -28,6 +28,31 @@ def reorganize_inout(json_obj, key):
         json_obj[key][index] = var_new
 
 
+def add_edges(json_obj):
+    json_obj['edges'] = []
+    label_incrementer = 0
+
+    for node_index in range(0, len(json_obj['node'])):
+        cur_node = json_obj['node'][node_index]
+
+        # input edges
+        for source in cur_node['input']:
+            json_obj['edges'].append({
+                'source': source,
+                'target': 'node_' + str(node_index),
+                'label': 'label_' + str(label_incrementer)
+            })
+            label_incrementer += 1
+
+        # output edge
+        json_obj['edges'].append({
+            'source': 'node_' + str(node_index),
+            'target': cur_node['output'][0],
+            'label': 'label_' + str(label_incrementer)
+        })
+        label_incrementer += 1
+
+
 def load_model(model_pb_path):
     model = onnx.load(model_pb_path)
     graph = model.graph
@@ -38,12 +63,14 @@ def load_model(model_pb_path):
     json_obj = json.loads(json_str)
     reorganize_inout(json_obj, 'input')
     reorganize_inout(json_obj, 'output')
+    add_edges(json_obj)
     return json.dumps(json_obj, sort_keys=True, indent=4, separators=(',', ': '))
 
 
 if __name__ == '__main__':
     import os
     import sys
     current_path = os.path.abspath(os.path.dirname(sys.argv[0]))
-    json_str = load_model(current_path + "/mock/inception_v1.pb")
+    # json_str = load_model(current_path + "/mock/inception_v1_model.pb")
+    json_str = load_model(current_path + "/mock/squeezenet_model.pb")
     print(json_str)

server/visualdl/mock/README.md

@@ -0,0 +1,8 @@
+## Note
+Test data squeezenet and inception_v1 are from onnx/models.
+
+https://github.com/onnx/models
+
+squeezenet has 53 inputs, 66 nodes and 1 output.
+
+inception_v1 has 117 inputs, 143 nodes and 1 output.

server/visualdl/mock/download_mock_models.sh

@@ -0,0 +1,20 @@
+
+# Download inception_v1 model
+curl -LOk https://s3.amazonaws.com/download.onnx/models/inception_v1.tar.gz
+
+tar -xvzf inception_v1.tar.gz
+cp inception_v1/model.pb inception_v1_model.pb
+
+rm -rf inception_v1
+rm inception_v1.tar.gz
+
+
+# Download squeezenet model
+curl -LOk https://s3.amazonaws.com/download.onnx/models/squeezenet.tar.gz
+
+tar -xvzf squeezenet.tar.gz
+cp squeezenet/model.pb squeezenet_model.pb
+
+rm -rf squeezenet
+rm squeezenet.tar.gz
+