Trimmed Mean Added to make throughput numbers more stable

tftrt/examples/benchmark_args.py

@@ -120,10 +120,21 @@ def __init__(self):
         self._parser.add_argument(
             "--num_warmup_iterations",
             type=int,
-            default=100,
+            default=200,
             help="Number of initial iterations skipped from timing."
         )
 
+        self._parser.add_argument(
+            "--trim_mean_percentage",
+            type=float,
+            default=0.1,
+            required=False,
+            help="Percentage used to trim step timing distribution from both "
+            "tails (fastest and slowest steps). 0.1 (default value) means that "
+            "10% of the fastest and slowest iteration will be removed for "
+            "model throughput computation."
+        )
+
         self._parser.add_argument(
             "--total_max_samples",
             type=int,

tftrt/examples/benchmark_runner.py

@@ -25,6 +25,8 @@
 from dataloading_utils import get_force_data_on_gpu_fn
 
 import numpy as np
+import scipy as sp
+import scipy.stats
 import tensorflow as tf
 
 from tensorflow.python.compiler.tensorrt import trt_convert as trt
@@ -500,11 +502,14 @@ def log_step(step_idx, display_every, iter_time, memcpyHtoD_time, dequeue_time):
 
             metrics['Total GPU Time (s)'] = int(np.ceil(np.sum(iter_times)))
             metrics['Throughput (samples/sec)'] = (
-                self._args.batch_size / np.mean(iter_times)
-            )
+                self._args.batch_size / sp.stats.trim_mean(
+                    iter_times, self._args.trim_mean_percentage))
 
             def timing_metrics(time_arr, log_prefix):
                 data = dict()
+                data[f"{log_prefix} Trim Mean [{self._args.trim_mean_percentage * 100}%] (ms)"] = (
+                    sp.stats.trim_mean(time_arr, self._args.trim_mean_percentage) * 1000
+                )
                 data[f"{log_prefix} 99th_percentile (ms)"] = np.percentile(
                     time_arr, q=99, interpolation='lower'
                 ) * 1000
@@ -522,9 +527,9 @@ def timing_metrics(time_arr, log_prefix):
 
             def log_value(key, val):
                 if isinstance(val, int):
-                    print(f"- {key:45s}: {val}")
+                    print(f"- {key:50s}: {val}")
                 else:
-                    print(f"- {key:45s}: {val:.2f}")
+                    print(f"- {key:50s}: {val:.2f}")
 
             for key, val in sorted(metrics.items()):
                 if isinstance(val, dict):