Update trt-engine-explorer to v0.1.6

Signed-off-by: Ilya Sherstyuk <isherstyuk@nvidia.com>

Author: ilyasher

tools/experimental/trt-engine-explorer/CHANGELOG.md

@@ -2,6 +2,11 @@
 
 Dates are in YYYY-MM-DD format.
 
+## v0.1.6 (2023-April)
+- Graph rendering:
+  - Add node highlighting option.
+  - Fix bug https://github.com/NVIDIA/TensorRT/issues/2779
+
 ## v0.1.5 (2022-12-06)
 - Updated requirements.txt for Ubuntu 20.04 and 22.04
 

tools/experimental/trt-engine-explorer/README.md

@@ -43,7 +43,7 @@ $ git clone https://github.com/NVIDIA/TensorRT.git
 ```
 
 ### 2. Create and activate a Python virtual environment
-The commands listed below create and activate a Python virtual enviornment named ```env_trex``` which is stored in a directory by the same name, and configures the current shell to use it as the default python environment.
+The commands listed below create and activate a Python virtual environment named ```env_trex``` which is stored in a directory by the same name, and configures the current shell to use it as the default python environment.
 
 ```
 $ cd TensorRT/tools/experimental/trt-engine-explorer

tools/experimental/trt-engine-explorer/setup.py

@@ -33,7 +33,7 @@ def main():
 
     setup(
         name="trex",
-        version="0.1.5",
+        version="0.1.6",
         description="TREX: TensorRT Engine Exploration Toolkit",
         long_description=open("README.md", "r", encoding="utf-8").read(),
         author="NVIDIA",

tools/experimental/trt-engine-explorer/trex/__init__.py

@@ -30,4 +30,4 @@
 from trex.compare_engines import *
 from trex.excel_summary import *
 
-__version__ = "0.1.5"
+__version__ = "0.1.6"

tools/experimental/trt-engine-explorer/trex/graphing.py

@@ -22,6 +22,7 @@
 
 import os
 import re
+import warnings
 from enum import Enum
 from graphviz import Digraph
 from typing import Callable, NamedTuple, List, Dict
@@ -434,6 +435,17 @@ def add_io(tensors: List):
     return label
 
 
+def layer_node_highlighter(node_id: str, highlighted_layers_ids: List[int]
+) -> Dict:
+    """Highlight a layer node.
+
+    Create a yellow hailo around the node.
+    """
+    should_highlight = highlighted_layers_ids and node_id in highlighted_layers_ids
+    formatting = {'penwidth': str(6), 'color': 'yellow'}
+    return formatting if should_highlight else {}
+
+
 def layer_node_configurable_renderer(
     layer: Layer,
     latency: float,
@@ -584,10 +596,12 @@ def handle_reformat(layer: Layer):
     except KeyError:
         layer_color = "#E5E7E9"
 
-    formatting = {'style': 'filled',
+    formatting = {'shape': 'Mrecord',
+                  'style': 'filled',
                   'tooltip': layer.tooltip(),
                   'fillcolor': layer_color,
-                  'color': 'white',}
+                  'color': 'lightgray',
+                  'fontname': 'Helvetica'}
     return formatting
 
 
@@ -623,11 +637,16 @@ def get_latency(plan: EnginePlan, layer: Layer, latency_type) -> float:
     return latency
 
 
+def get_dot_id(layer_name: str) -> str:
+    return layer_name.replace(":", "###") # f"l_{dot_node_id}"
+
+
 class DotGraph(object):
     """This class converts a TensorRT plan into Graphviz DOT graphs"""
     def __init__(self,
         plan: EnginePlan,
         layer_node_formatter: Callable,
+        layer_node_highlighter: Callable=layer_node_highlighter,
         layer_node_renderer: Callable=layer_node_configurable_renderer,
         region_formatter: Callable=region_precision_formatter,
         display_layer_names: bool=True,
@@ -642,11 +661,13 @@ def __init__(self,
         display_region_names: bool=False,
         display_edge_name: bool=False,
         display_edge_details: bool=True,
+        highlight_layers: list=None,
     ):
         plan_graph = PlanGraph(
             plan, display_regions, display_constants, display_forking_regions)
         self.dot = Digraph()
         self.layer_node_formatter = layer_node_formatter
+        self.layer_node_highlighter = layer_node_highlighter
         self.layer_node_renderer = layer_node_renderer
         self.region_formatter = region_formatter
         self.expand_layer_details = expand_layer_details
@@ -659,6 +680,14 @@ def __init__(self,
         self.display_region_names = display_region_names
         self.display_edge_name = display_edge_name
         self.display_edge_details = display_edge_details
+        # Get the node names of the layers to highlight
+        self.highlighted_layers_ids = None
+        if highlight_layers:
+            try:
+                highlight_layers_name = plan.df['Name'].iloc[highlight_layers].to_list()
+                self.highlighted_layers_ids = [get_dot_id(name) for name in highlight_layers_name]
+            except IndexError:
+                warnings.warn("The layers indices specified for highlighting are incorrect")
 
         node_name_2_node_id = {}
         self.__add_dot_region_nodes(plan_graph, node_name_2_node_id)
@@ -672,7 +701,7 @@ def __init__(self,
     def __add_dot_region_nodes(self, plan_graph, node_name_2_node_id):
         dot_node_id = 0
         for mem_node in plan_graph.memory_nodes:
-            node_name_2_node_id[mem_node.name] = dot_id = mem_node.name.replace(":", "###") #f"r_{dot_node_id}"
+            node_name_2_node_id[mem_node.name] = dot_id = get_dot_id(mem_node.name)
             self.__create_dot_region_node(dot_id, mem_node.tensor, mem_node.is_user, mem_node.region_gen)
             dot_node_id += 1
 
@@ -681,7 +710,7 @@ def __add_dot_layer_nodes(self, plan, plan_graph, node_name_2_node_id):
             layer = layer_node.layer
             latency = get_latency(plan, layer, self.latency_type)
             if not layer.type == 'Constant' or plan_graph.include_constants:
-                dot_id = layer.name.replace(":", "###") # f"l_{dot_node_id}"
+                dot_id = get_dot_id(layer.name)
                 node_name_2_node_id[layer.name] = dot_id
                 self.__create_dot_layer_node(
                     dot_id, layer, latency, layer_node_renderer=self.layer_node_renderer)
@@ -713,6 +742,7 @@ def __create_dot_layer_node(
         self, node_id: int, layer: Layer, latency: float, layer_node_renderer: Callable
     ):
         formatting = self.layer_node_formatter(layer)
+        formatting.update(self.layer_node_highlighter(node_id, self.highlighted_layers_ids))
         self.dot.node(
             str(node_id),
             layer_node_renderer(
@@ -721,8 +751,7 @@ def __create_dot_layer_node(
                 expand_layer_details=self.expand_layer_details,
                 display_layer_names=self.display_layer_names,
                 stack_layer_names=self.stack_layer_names),
-            shape='Mrecord',
-            fontname="Helvetica", **formatting)
+                **formatting)
 
     def __create_dot_edge(self, src, end, tensor, region_gen):
         def generation_color(gen: int, line_color: str) -> str:

tools/experimental/trt-engine-explorer/trex/parser.py

@@ -172,7 +172,23 @@ def convert_deconv(raw_layers: List) -> List:
                 pass
         return raw_layers
 
+    def fix_metadata(raw_layers: List) -> List:
+        """TensorRT 8.6 introduced the Metadata field, with a non-ASCII character
+        that triggers an SVG rendering error. This function replaces this character.
+
+        See: https://github.com/NVIDIA/TensorRT/issues/2779
+        """
+        TRT_METADATA_DELIM = '\x1E'
+        for l in raw_layers:
+            try:
+                if TRT_METADATA_DELIM in l['Metadata']:
+                    l['Metadata'] = l['Metadata'].replace(TRT_METADATA_DELIM, '+')
+            except KeyError:
+                pass
+        return raw_layers
+
     raw_layers, bindings = read_graph_file(graph_file)
+    raw_layers = fix_metadata(raw_layers)
     raw_layers = convert_deconv(raw_layers)
     raw_layers = disambiguate_layer_names(raw_layers)
     raw_layers, bindings = filter_profiles(raw_layers, bindings, profile_id)