#!/usr/bin/env python
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""This tool creates an html visualization of a TensorFlow Lite graph.
Example usage:
python visualize.py foo.tflite foo.html
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import json
import os
import sys
from tensorflow.python.platform import resource_loader
# Schema to use for flatbuffers
_SCHEMA = "third_party/tensorflow/contrib/lite/schema/schema.fbs"
# TODO(angerson): fix later when rules are simplified..
_SCHEMA = resource_loader.get_path_to_datafile("../schema/schema.fbs")
_BINARY = resource_loader.get_path_to_datafile("../../../../flatbuffers/flatc")
# Account for different package positioning internal vs. external.
if not os.path.exists(_BINARY):
_BINARY = resource_loader.get_path_to_datafile(
"../../../../../flatbuffers/flatc")
if not os.path.exists(_SCHEMA):
raise RuntimeError("Sorry, schema file cannot be found at %r" % _SCHEMA)
if not os.path.exists(_BINARY):
raise RuntimeError("Sorry, flatc is not available at %r" % _BINARY)
# A CSS description for making the visualizer
_CSS = """
"""
_D3_HTML_TEMPLATE = """
"""
class OpCodeMapper(object):
"""Maps an opcode index to an op name."""
def __init__(self, data):
self.code_to_name = {}
for idx, d in enumerate(data["operator_codes"]):
self.code_to_name[idx] = d["builtin_code"]
def __call__(self, x):
if x not in self.code_to_name:
s = ""
else:
s = self.code_to_name[x]
return "%s (opcode=%d)" % (s, x)
class DataSizeMapper(object):
"""For buffers, report the number of bytes."""
def __call__(self, x):
if x is not None:
return "%d bytes" % len(x)
else:
return "--"
class TensorMapper(object):
"""Maps a list of tensor indices to a tooltip hoverable indicator of more."""
def __init__(self, subgraph_data):
self.data = subgraph_data
def __call__(self, x):
html = ""
html += ""
for i in x:
tensor = self.data["tensors"][i]
html += str(i) + " "
html += tensor["name"] + " "
html += str(tensor["type"]) + " "
html += (repr(tensor["shape"]) if "shape" in tensor else "[]") + " "
html += ""
html += repr(x)
html += ""
return html
def GenerateGraph(subgraph_idx, g, opcode_mapper):
"""Produces the HTML required to have a d3 visualization of the dag."""
def TensorName(idx):
return "t%d" % idx
def OpName(idx):
return "o%d" % idx
edges = []
nodes = []
first = {}
pixel_mult = 50 # TODO(aselle): multiplier for initial placement
for op_index, op in enumerate(g["operators"]):
for tensor_input_position, tensor_index in enumerate(op["inputs"]):
if tensor_index not in first:
first[tensor_index] = (
op_index * pixel_mult,
tensor_input_position * pixel_mult - pixel_mult / 2)
edges.append({
"source": TensorName(tensor_index),
"target": OpName(op_index)
})
for tensor_index in op["outputs"]:
edges.append({
"target": TensorName(tensor_index),
"source": OpName(op_index)
})
nodes.append({
"id": OpName(op_index),
"name": opcode_mapper(op["opcode_index"]),
"group": 2,
"x": pixel_mult,
"y": op_index * pixel_mult
})
for tensor_index, tensor in enumerate(g["tensors"]):
initial_y = (
first[tensor_index] if tensor_index in first else len(g["operators"]))
nodes.append({
"id": TensorName(tensor_index),
"name": "%s (%d)" % (tensor["name"], tensor_index),
"group": 1,
"x": 2,
"y": initial_y
})
graph_str = json.dumps({"nodes": nodes, "edges": edges})
html = _D3_HTML_TEMPLATE % (graph_str, subgraph_idx)
return html
def GenerateTableHtml(items, keys_to_print, display_index=True):
"""Given a list of object values and keys to print, make an HTML table.
Args:
items: Items to print an array of dicts.
keys_to_print: (key, display_fn). `key` is a key in the object. i.e.
items[0][key] should exist. display_fn is the mapping function on display.
i.e. the displayed html cell will have the string returned by
`mapping_fn(items[0][key])`.
display_index: add a column which is the index of each row in `items`.
Returns:
An html table.
"""
html = ""
# Print the list of items
html += "
\n"
html += "
\n"
if display_index:
html += "
index
"
for h, mapper in keys_to_print:
html += "
%s
" % h
html += "
\n"
for idx, tensor in enumerate(items):
html += "
\n"
if display_index:
html += "
%d
" % idx
# print tensor.keys()
for h, mapper in keys_to_print:
val = tensor[h] if h in tensor else None
val = val if mapper is None else mapper(val)
html += "
%s
\n" % val
html += "
\n"
html += "
\n"
return html
def CreateHtmlFile(tflite_input, html_output):
"""Given a tflite model in `tflite_input` file, produce html description."""
# Convert the model into a JSON flatbuffer using flatc (build if doesn't
# exist.
if not os.path.exists(tflite_input):
raise RuntimeError("Invalid filename %r" % tflite_input)
if tflite_input.endswith(".tflite") or tflite_input.endswith(".bin"):
# Run convert
cmd = (
_BINARY + " -t "
"--strict-json --defaults-json -o /tmp {schema} -- {input}".format(
input=tflite_input, schema=_SCHEMA))
print(cmd)
os.system(cmd)
real_output = ("/tmp/" + os.path.splitext(
os.path.split(tflite_input)[-1])[0] + ".json")
data = json.load(open(real_output))
elif tflite_input.endswith(".json"):
data = json.load(open(tflite_input))
else:
raise RuntimeError("Input file was not .tflite or .json")
html = ""
html += _CSS
html += "
TensorFlow Lite Model"
data["filename"] = tflite_input # Avoid special case
toplevel_stuff = [("filename", None), ("version", None), ("description",
None)]
html += "
\n"
for key, mapping in toplevel_stuff:
if not mapping:
mapping = lambda x: x
html += "
%s
%s
\n" % (key, mapping(data.get(key)))
html += "
\n"
# Spec on what keys to display
buffer_keys_to_display = [("data", DataSizeMapper())]
operator_keys_to_display = [("builtin_code", None)]
for subgraph_idx, g in enumerate(data["subgraphs"]):
# Subgraph local specs on what to display
html += "