提交 7aee3b6e 编写于 作者: D daminglu 提交者: GitHub

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
上级 4eda2386
......@@ -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,6 +63,7 @@ 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=(',', ': '))
......@@ -45,5 +71,6 @@ 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)
## 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.
\ No newline at end of file
# 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
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册