Skip to content

Commit 962fb48

Browse files
committed
minor update
1 parent 8c69f8d commit 962fb48

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

tools/perf/perf_run.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,7 @@ def run_dynamo(model, input_tensors, params, precision, batch_size):
282282
cache_built_engines=params.get("cache_built_engines", False),
283283
reuse_cached_engines=params.get("reuse_cached_engines", False),
284284
use_python_runtime=params.get("use_python_runtime", False),
285-
optimization_level=5,
285+
optimization_level=params.get("optimization_level", 5),
286286
)
287287
end_compile = timeit.default_timer()
288288
compile_time_s = end_compile - start_compile
@@ -455,7 +455,7 @@ def run_tensorrt(
455455
config = builder.create_builder_config()
456456
if precision == "fp16":
457457
config.set_flag(trt.BuilderFlag.FP16)
458-
config.builder_optimization_level = 5
458+
config.builder_optimization_level = params.get("optimization_level", 5)
459459
start_compile = timeit.default_timer()
460460
serialized_engine = builder.build_serialized_network(network, config)
461461
end_compile = timeit.default_timer()
@@ -680,6 +680,12 @@ def run(
680680
action="store_true",
681681
help="Truncate long and double weights in the network in Torch-TensorRT",
682682
)
683+
arg_parser.add_argument(
684+
"--optimization_level",
685+
type=int,
686+
default=5,
687+
help="Builder optimization level for TensorRT",
688+
)
683689
arg_parser.add_argument(
684690
"--is_trt_engine",
685691
action="store_true",

0 commit comments

Comments
 (0)