Adding flag for number of batches to use for TF-TRT engine building

tensorflow · Mar 10, 2022 · dd03467 · dd03467
1 parent 920597b
commit dd03467
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/tftrt/examples/benchmark_args.py b/tftrt/examples/benchmark_args.py
@@ -180,6 +180,15 @@ def __init__(self):
             help="Minimum number of TensorFlow ops in a TRT engine."
         )
 
+        self._parser.add_argument(
+            "--num_build_batches",
+            type=int,
+            default=1,
+            help="How many iterations(batches) to use to build the TF-TRT "
+            "engines. If not supplied, only one batch will be used. This "
+            "parameter has only an effect if `--optimize_offline=True`"
+        )
+
         self._parser.add_argument(
             "--num_calib_batches",
             type=int,

diff --git a/tftrt/examples/benchmark_runner.py b/tftrt/examples/benchmark_runner.py
@@ -228,7 +228,8 @@ def engine_build_input_fn(num_batches, model_phase):
             if self._args.optimize_offline:
 
                 offline_opt_input_fn = lambda: engine_build_input_fn(
-                    num_batches=1, model_phase="Building"
+                    num_batches=self._args.num_build_batches,
+                    model_phase="Building"
                 )
 
                 with timed_section("Building TensorRT engines"):