提交 db44c87e 编写于 作者: Q qiaolongfei

init graph

上级 eeaa1bd4
import onnx
import json
from google.protobuf.json_format import MessageToJson
def reorganize_inout(json_obj, key):
"""
:param json_obj: the model's json obj
:param key: "input or output"
:return:
"""
for index in range(len(json_obj[key])):
input = json_obj[key][index]
input_new = dict()
# set name
input_new['name'] = input['name']
tensor_type = input['type']['tensorType']
# set data_type
input_new['data_type'] = tensor_type['elemType']
# set shape
shape = [dim['dimValue'] for dim in tensor_type['shape']['dim']]
input_new['shape'] = shape
json_obj[key][index] = input_new
def load_model(model_pb_path):
model = onnx.load(model_pb_path)
graph = model.graph
del graph.initializer[:]
# to json string
json_str = MessageToJson(model.graph)
json_obj = json.loads(json_str)
reorganize_inout(json_obj, 'input')
reorganize_inout(json_obj, 'output')
return json.dumps(json_obj, sort_keys=True, indent=4, separators=(',', ': '))
json_str = load_model("/Users/qiaolongfei/project/onnx/model/inception_v1/model.pb")
print(json_str)
...@@ -14,6 +14,7 @@ from visualdl.log import logger ...@@ -14,6 +14,7 @@ from visualdl.log import logger
import visualdl.mock.data as mock_data import visualdl.mock.data as mock_data
import visualdl.mock.tags as mock_tags import visualdl.mock.tags as mock_tags
import storage import storage
import graph
app = Flask(__name__, static_url_path="") app = Flask(__name__, static_url_path="")
...@@ -135,6 +136,12 @@ def scalars(): ...@@ -135,6 +136,12 @@ def scalars():
return Response(json.dumps(result), mimetype='application/json') return Response(json.dumps(result), mimetype='application/json')
@app.route('/data/plugin/graphs/graph')
def graph():
model_json = graph.load_model("")
return Response(model_json, mimetype='application/json')
if __name__ == '__main__': if __name__ == '__main__':
logger.info(" port=" + str(options.port)) logger.info(" port=" + str(options.port))
app.run(debug=False, host=options.host, port=options.port) app.run(debug=False, host=options.host, port=options.port)
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册