Adding JSON export

Author: DEKHTIARJonathan

tftrt/examples/benchmark_args.py

@@ -237,6 +237,14 @@ def __init__(self):
 
         # =========================== DEBUG Flags ========================== #
 
+        self._parser.add_argument(
+            "--export_metrics_json_path",
+            type=str,
+            default=None,
+            help="If set, the script will export runtime metrics and arguments "
+            "to the set location in JSON format for further processing."
+        )
+
         self._add_bool_argument(
             name="debug",
             default=False,

tftrt/examples/benchmark_runner.py

@@ -5,6 +5,8 @@
 import os
 
 import abc
+import copy
+import json
 import logging
 import sys
 import time
@@ -104,6 +106,31 @@ def _debug_print(self, msg):
         if self._args.debug:
             print(f"[DEBUG] {msg}")
 
+    def _export_runtime_metrics_to_json(self, metric_dict):
+
+        metric_dict = {
+          # Creating a copy to avoid modifying the original
+          "results": copy.deepcopy(metric_dict),
+          "runtime_arguments": vars(self._args)
+        }
+
+        json_path = self._args.export_metrics_json_path
+        if json_path is not None:
+            try:
+                with open(json_path, 'w') as json_f:
+                    json_string = json.dumps(
+                        metric_dict,
+                        default=lambda o: o.__dict__,
+                        sort_keys=True,
+                        indent=4
+                    )
+                    print(json_string, file=json_f)
+            except Exception as e:
+                print(
+                    "[ERROR] Impossible to save JSON File at path: "
+                    f"{json_path}.\nError: {str(e)}"
+                )
+
     def _get_graph_func(self):
         """Retreives a frozen SavedModel and applies TF-TRT
         use_tftrt: bool, if true use TensorRT
@@ -389,16 +416,15 @@ def log_step(step_idx, display_every, iter_time):
 
         with timed_section("Metric Computation"):
 
+            metrics = dict()
+
             if not self._args.use_synthetic_data:
                 metric, metric_units = self.evaluate_model(
                     data_aggregator.predicted_dict,
                     data_aggregator.expected_dict, bypass_data_to_eval
                 )
-                print(f"- {metric_units:35s}: {metric:.2f}")
+                metrics["Metric"] = {metric_units: metric}
 
-            metrics = dict()
-
-            if not self._args.use_synthetic_data:
                 metrics["Total Samples Processed"] = (
                     data_aggregator.total_samples_processed
                 )
@@ -419,10 +445,18 @@ def log_step(step_idx, display_every, iter_time):
             metrics['GPU Latency Min (ms)'] = np.min(run_times) * 1000
             metrics['GPU Latency Max (ms)'] = np.max(run_times) * 1000
 
-            for key, val in sorted(metrics.items()):
+            self._export_runtime_metrics_to_json(metrics)
+
+            def log_value(key, val):
                 if isinstance(val, int):
                     print(f"- {key:35s}: {val}")
                 else:
                     print(f"- {key:35s}: {val:.2f}")
 
+            for key, val in sorted(metrics.items()):
+                if isinstance(val, dict):
+                    log_value(*list(val.items())[0])
+                else:
+                    log_value(key, val)
+
         print()  # visual spacing