提交 3f63067c 编写于 作者: S superjom

Merge branch 'develop' of github.com:PaddlePaddle/VisualDL into feature/add_image_component

...@@ -28,6 +28,31 @@ def reorganize_inout(json_obj, key): ...@@ -28,6 +28,31 @@ def reorganize_inout(json_obj, key):
json_obj[key][index] = var_new 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): def load_model(model_pb_path):
model = onnx.load(model_pb_path) model = onnx.load(model_pb_path)
graph = model.graph graph = model.graph
...@@ -38,6 +63,7 @@ def load_model(model_pb_path): ...@@ -38,6 +63,7 @@ def load_model(model_pb_path):
json_obj = json.loads(json_str) json_obj = json.loads(json_str)
reorganize_inout(json_obj, 'input') reorganize_inout(json_obj, 'input')
reorganize_inout(json_obj, 'output') reorganize_inout(json_obj, 'output')
add_edges(json_obj)
return json.dumps(json_obj, sort_keys=True, indent=4, separators=(',', ': ')) return json.dumps(json_obj, sort_keys=True, indent=4, separators=(',', ': '))
...@@ -45,5 +71,6 @@ if __name__ == '__main__': ...@@ -45,5 +71,6 @@ if __name__ == '__main__':
import os import os
import sys import sys
current_path = os.path.abspath(os.path.dirname(sys.argv[0])) 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) 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 http://visualdl.bj.bcebos.com/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 http://visualdl.bj.bcebos.com/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.
先完成此消息的编辑!
想要评论请 注册