Integrate trex v0.1.2

Signed-off-by: Rajeev Rao <[email protected]>

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

@@ -2,7 +2,7 @@
 
 Dates are in YYYY-MM-DD format.
 
-## v0.1.1 (2022-06-01)
+## v0.1.2 (2022-06-01)
 - Graph rendering:
   - Add timing information to engine graph nodes.
   - Change layer coloring scheme.

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

@@ -26,7 +26,7 @@ When `trtexec` times individual layers, the total engine latency (computed by su
 
 To measure per-layer execution times, when `trtexec` enqueues kernel layers for execution in a stream, it places CUDA event objects between the layers to monitor the start and completion of each layer. These CUDA events add a small overhead which is more noticeable with smaller networks (shallow and narrow networks or networks with small activation data).
 
-## Supported TenorRT Versions
+## Supported TensorRT Versions
 Starting with TensorRT 8.2, engine-plan graph and profiling data can be exported to JSON files. `trex` supports TensortRT 8.2 and 8.4.
 
 `trex` has only been tested on Ubuntu 18.04 LTS, with Python 3.6.<br>

tools/experimental/trt-engine-explorer/examples/pytorch/resnet/README.md

@@ -9,7 +9,7 @@ The example uses PyTorch,  TorchVision and Nvidia's [PyTorch Quantization Toolki
 python3 -m pip install -r requirements.txt
 ```
 
-### Introduction
+### Description
 This example walks you through the process of optimizing a TensorRT engine created from QAT ResNet18.
 File `resnet_example.py` contains code to generate several variations of QAT ResNet18, each one performing better than the previous.
 
@@ -25,4 +25,4 @@ File `resnet_example.py` contains code to generate several variations of QAT Res
     ./trt-engine-explorer-OSS/examples/pytorch/resnet/process_resnet.sh
     ```
 
-* Finally, review the results using the `example_qat_resnet18.ipynb` notebook.
+* Finally, review the results using the `example_qat_resnet18.ipynb` notebook.

tools/experimental/trt-engine-explorer/examples/pytorch/resnet/example_qat_resnet18.ipynb

@@ -72,7 +72,7 @@
     "        f'{engine_path}.graph.json',\n",
     "        f'{engine_path}.profile.json',\n",
     "        f\"{engine_path}.profile.metadata.json\",\n",
-    "        engine_name)\n",
+    "        name = engine_name)\n",
     "    return plan\n",
     "\n",
     "\n",
@@ -159,7 +159,7 @@
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -173,7 +173,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.9"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

tools/experimental/trt-engine-explorer/examples/pytorch/resnet/requirements.txt

@@ -1,6 +1,6 @@
 --find-links https://download.pytorch.org/whl/cu113
 
-torch
-torchvision
 nvidia-pyindex
 pytorch-quantization --extra-index-url https://pypi.ngc.nvidia.com
+torch
+torchvision

tools/experimental/trt-engine-explorer/install.sh

@@ -16,6 +16,16 @@
 # limitations under the License.
 #
 
+#
+# TREx installation script.
+#
+# Usage:
+#   $ source install.sh
+#
+
+
+sudo apt install graphviz
+sudo apt install virtualenv
 python3 -m virtualenv env_trex
 source env_trex/bin/activate
 python3 -m pip install -e .

tools/experimental/trt-engine-explorer/notebooks/compare_engines.ipynb

@@ -134,7 +134,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": false
+   },
    "outputs": [],
    "source": [
     "compare_engines_layer_details(plans[0], plans[1])"
@@ -164,6 +166,43 @@
     "print(plan2.name)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "report_card_table_view(plan1);\n",
+    "print(plan1.name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "report_card_table_view(plan2);\n",
+    "print(plan2.name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "for plan in plans:\n",
+    "    graph = to_dot(plan, layer_type_formatter, display_regions=True, expand_layer_details=True)\n",
+    "    render_dot(graph, plan.name, 'svg')"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -224,7 +263,7 @@
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -238,7 +277,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.9"
+   "version": "3.8.10"
   }
  },
  "nbformat": 4,

tools/experimental/trt-engine-explorer/notebooks/engine_report_card.ipynb

@@ -94,7 +94,7 @@
    "source": [
     "if fc.selected is not None:\n",
     "    engine_name = fc.selected\n",
-    "    \n",
+    "\n",
     "assert engine_name is not None\n",
     "plan = EnginePlan(f\"{engine_name}.graph.json\", f\"{engine_name}.profile.json\", f\"{engine_name}.profile.metadata.json\")"
    ]

tools/experimental/trt-engine-explorer/notebooks/tutorial.ipynb

@@ -34,7 +34,7 @@
    },
    "outputs": [],
    "source": [
-    "#! python3 utils/process_engine.py tests/inputs/mobilenet.qat.onnx test/inputs best"
+    "# !python3 ../utils/process_engine.py ../tests/inputs/mobilenet.qat.onnx ../tests/inputs best"
    ]
   },
   {

tools/experimental/trt-engine-explorer/requirements.txt

@@ -1,7 +1,9 @@
 setuptools # for qgrid
 wheel # for qgrid
+protobuf==3.16.0
+onnx==1.11.0
 numpy
-pandas
+pandas==1.1.5
 plotly
 qgrid
 graphviz
@@ -13,5 +15,4 @@ ipyfilechooser
 jupyterlab
 jupyter-dash
 pytest
-dtale
-onnx
+dtale==2.2.0

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

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

tools/experimental/trt-engine-explorer/tests/test_basic.py

@@ -48,7 +48,7 @@ def test_summary(plan):
     assert len(plan.layers) == 72
     assert plan.total_weights_size == 3469760
     assert plan.total_act_size == 16629344
-    d = plan.summary_dict()
+    d = summary_dict(plan)
     assert d["Inputs"] == "[input1: [1, 3, 224, 224]xFP32 NCHW]"
     assert d["Average time"] == "0.470 ms"
     assert d["Layers"] == "72"

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

@@ -29,4 +29,4 @@
 from trex.report_card import *
 from trex.compare_engines import *
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"

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

@@ -21,6 +21,8 @@
 
 
 import numpy as np
+from typing import Dict
+import pandas as pd
 
 
 # This dictionary compresses JSON's long format description strings.
@@ -59,7 +61,7 @@
 
 class Activation:
     """Convenience class wrapping activation regions."""
-    def __init__(self, raw_dict):
+    def __init__(self, raw_dict: Dict):
         def parse_tensor_info(desc):
             if 'Int8' in desc:
                 precision = 'INT8'
@@ -76,14 +78,17 @@ def parse_tensor_info(desc):
             elif 'Bool' in desc:
                 precision = 'BOOL'
                 data_size = 4
+            elif desc == "Unknown format":
+                precision = 'Unknown'
+                data_size = 0
             else:
                 raise ValueError(f"Uknown precision {desc}")
             return precision, data_size
 
         self.name = raw_dict['Name']
         self.shape = raw_dict['Dimensions']
         format = raw_dict['Format/Datatype'].replace('.', '')
-        self.format = _regionFormatDict[format]
+        self.format = _regionFormatDict.get(format,"Unknown format")
         self.precision, self.data_size = parse_tensor_info(self.format)
         self.size_bytes = np.prod(self.shape) * self.data_size
 
@@ -97,7 +102,8 @@ def tooltip(self):
     def __repr__(self):
         return f"{self.name}: {str(self.shape)}x{self.format}"
 
-def create_activations(layer):
+
+def create_activations(layer: pd.Series):
     inputs = [Activation(tensor) for tensor in layer.Inputs]
     outputs = [Activation(tensor) for tensor in layer.Outputs]
     return inputs, outputs

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

@@ -31,6 +31,7 @@
 from .interactive import *
 from .misc import stack_dicts
 from .activations import create_activations
+from .engine_plan import summary_dict
 
 
 def get_plans_names(plans: List[EnginePlan]):
@@ -201,8 +202,18 @@ def latency_per_type(title):
 
 def compare_engines_summaries_tbl(plans: List[EnginePlan], orientation: str='vertical'):
     """Display a tabular comparison of several engine plans."""
-    summary_dicts = [plan.summary_dict() for plan in plans]
-    merged_summaries = stack_dicts(summary_dicts, empty_placeholder="")
+
+    merged_summaries = {}
+    summary_dicts_list = (
+        [summary_dict(plan) for plan in plans],
+        [plan.performance_summary for plan in plans],
+        [plan.device_properties for plan in plans],
+        [plan.builder_cfg for plan in plans]
+    )
+
+    for d in summary_dicts_list:
+        merged_summaries.update(stack_dicts(d, empty_placeholder=""))
+
     if orientation == 'vertical':
         df = pd.DataFrame.from_dict(
             merged_summaries, orient='index', columns=get_plans_names(plans))

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

@@ -114,35 +114,9 @@ def compute_summary(self):
         compute_summary(self)
         self.device_properties = get_device_properties(profiling_metadata_file)
         self.performance_summary = get_performance_summary(profiling_metadata_file)
+        self.builder_cfg = get_builder_config(build_metadata_file)
         assert self._df is not None, f"Failed parsing plan file {graph_file}"
 
-    def summary_dict(self):
-        """Create a dictionary of important attributes of the engine plan."""
-        MB_1 = 1024 * 1024
-        bindings = self.get_bindings()
-
-        d = {
-            "Inputs": f"{bindings[0]}",
-            "Average time": f"{self.total_runtime:.3f} ms",
-            "Layers": f"{len(self.df)}",
-            "Weights": f"{self.total_weights_size / MB_1 :.1f} MB",
-            "Activations": f"{self.total_act_size/ MB_1 :.1f} MB",
-        }
-        d.update(self.device_properties)
-        d.update(self.performance_summary)
-        return d
-
-    def print_summary(self):
-        d = self.summary_dict()
-        for k,v in d.items():
-            print(f"{k}: {v}")
-
-    def summary(self):
-        return self.print_summary()
-
-    def as_dataframe(self):
-        return self._df
-
     @property
     def df(self):
         return self._df
@@ -162,3 +136,33 @@ def get_bindings(self) -> Tuple[List[Activation], List[Activation]]:
             inputs += [inp for inp in layer.inputs if inp.name in self.bindings]
             outputs += [outp for outp in layer.outputs if outp.name in self.bindings]
         return inputs, outputs
+
+    def summary(self):
+        return print_summary(self)
+
+
+def summary_dict(plan: EnginePlan):
+    """Create a dictionary of important attributes of the engine plan."""
+    MB_1 = 1024 * 1024
+    bindings = plan.get_bindings()
+
+    d = {
+        "Inputs": f"{bindings[0]}",
+        "Average time": f"{plan.total_runtime:.3f} ms",
+        "Layers": f"{len(plan.df)}",
+        "Weights": f"{plan.total_weights_size / MB_1 :.1f} MB",
+        "Activations": f"{plan.total_act_size/ MB_1 :.1f} MB",
+    }
+    return d
+
+
+def print_summary(plan: EnginePlan):
+    def print_dict(d: Dict):
+        for k,v in d.items():
+            print(f"\t{k}: {v}")
+    print("Model:")
+    print_dict(summary_dict(plan))
+    print("Device Properties:")
+    print_dict(plan.device_properties)
+    print("Performance Summary:")
+    print_dict(plan.performance_summary)

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

@@ -69,7 +69,7 @@ def __init__(self, plan):
 
         # Create region generations.
 
-        # Each time we switch between C events to P events we create a new
+        # Each time we switch between C events and P events we create a new
         # Region generation.
         regions_gens = {}
         for region, region_evts in story.items():

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

@@ -59,8 +59,11 @@ def display_df_dtale(
     precision: int=4):
     """Display a Pandas dataframe using a dtale widget"""
     d = dtale.show(
-        df, drop_index=True, allow_cell_edits=False,
-        precision=precision, nan_display=nan_display)
+        df,
+        drop_index=True,
+        allow_cell_edits=False,
+        precision=precision,
+        nan_display=nan_display)
     if range_highlights is not None:
         d.update_settings(range_highlights=range_highlights, background_mode='range')
     display(d)