| 26 | self._proto = proto |
| 27 | |
| 28 | def render_html(self): |
| 29 | json_obj = { |
| 30 | "nodes": [], |
| 31 | "links": [] |
| 32 | } |
| 33 | |
| 34 | json_printer = _Printer() |
| 35 | |
| 36 | for op in self._proto.op: |
| 37 | op_json = json_printer._MessageToJsonObject(op) |
| 38 | op_json["id"] = op_json["name"] |
| 39 | op_json["node_type"] = "op" |
| 40 | json_obj["nodes"].append(op_json) |
| 41 | |
| 42 | for tensor in self._proto.tensors: |
| 43 | tensor_json = json_printer._MessageToJsonObject(tensor) |
| 44 | |
| 45 | tensor_json["id"] = tensor_json["name"] |
| 46 | if "floatData" in tensor_json and \ |
| 47 | len(tensor_json["floatData"]) > THREASHOLD: |
| 48 | del tensor_json["floatData"] |
| 49 | if "int32Data" in tensor_json and \ |
| 50 | len(tensor_json["int32Data"]) > THREASHOLD: |
| 51 | del tensor_json["int32Data"] |
| 52 | tensor_json["node_type"] = "tensor" |
| 53 | json_obj["nodes"].append(tensor_json) |
| 54 | |
| 55 | node_ids = [node["id"] for node in json_obj["nodes"]] |
| 56 | |
| 57 | tensor_to_op = {} |
| 58 | for op in self._proto.op: |
| 59 | for tensor in op.output: |
| 60 | tensor_to_op[tensor] = op.name |
| 61 | |
| 62 | for op in json_obj["nodes"]: |
| 63 | if "input" in op: |
| 64 | for input in op["input"]: |
| 65 | if input in node_ids and op["name"] in node_ids: |
| 66 | # for weights |
| 67 | json_obj["links"].append( |
| 68 | {"source": input, "target": op["name"]}) |
| 69 | elif input in tensor_to_op and \ |
| 70 | tensor_to_op[input] in node_ids: |
| 71 | # for intermediate tensor |
| 72 | json_obj["links"].append( |
| 73 | {"source": tensor_to_op[input], |
| 74 | "target": op["name"]}) |
| 75 | else: |
| 76 | # for input |
| 77 | json_obj["nodes"].append({ |
| 78 | "id": input, |
| 79 | "name": input, |
| 80 | "node_type": "input" |
| 81 | }) |
| 82 | json_obj["links"].append( |
| 83 | {"source": input, "target": op["name"]}) |
| 84 | |
| 85 | json_msg = json.dumps(json_obj, cls=NPEncoder) |