Display the names of the custom ops in the graph display, as well "Ops" and "Operator Codes" tables.
PiperOrigin-RevId: 356857986
Change-Id: I82c835c572877fedd7ce1a452a83d25819407d9f
diff --git a/tensorflow/lite/tools/visualize.py b/tensorflow/lite/tools/visualize.py
index 92c51f2..e51f455 100644
--- a/tensorflow/lite/tools/visualize.py
+++ b/tensorflow/lite/tools/visualize.py
@@ -234,6 +234,8 @@
self.code_to_name = {}
for idx, d in enumerate(data["operator_codes"]):
self.code_to_name[idx] = BuiltinCodeToName(d["builtin_code"])
+ if self.code_to_name[idx] == "CUSTOM":
+ self.code_to_name[idx] = NameListToString(d["custom_code"])
def __call__(self, x):
if x not in self.code_to_name:
@@ -450,7 +452,7 @@
# Spec on what keys to display
buffer_keys_to_display = [("data", DataSizeMapper())]
operator_keys_to_display = [("builtin_code", BuiltinCodeToName),
- ("custom_code", None),
+ ("custom_code", NameListToString),
("version", None)]
# Update builtin code fields.
