@@ -282,7 +282,7 @@ def run_dynamo(model, input_tensors, params, precision, batch_size):
282
282
cache_built_engines = params .get ("cache_built_engines" , False ),
283
283
reuse_cached_engines = params .get ("reuse_cached_engines" , False ),
284
284
use_python_runtime = params .get ("use_python_runtime" , False ),
285
- optimization_level = 5 ,
285
+ optimization_level = params . get ( "optimization_level" , 5 ) ,
286
286
)
287
287
end_compile = timeit .default_timer ()
288
288
compile_time_s = end_compile - start_compile
@@ -455,7 +455,7 @@ def run_tensorrt(
455
455
config = builder .create_builder_config ()
456
456
if precision == "fp16" :
457
457
config .set_flag (trt .BuilderFlag .FP16 )
458
- config .builder_optimization_level = 5
458
+ config .builder_optimization_level = params . get ( "optimization_level" , 5 )
459
459
start_compile = timeit .default_timer ()
460
460
serialized_engine = builder .build_serialized_network (network , config )
461
461
end_compile = timeit .default_timer ()
@@ -680,6 +680,12 @@ def run(
680
680
action = "store_true" ,
681
681
help = "Truncate long and double weights in the network in Torch-TensorRT" ,
682
682
)
683
+ arg_parser .add_argument (
684
+ "--optimization_level" ,
685
+ type = int ,
686
+ default = 5 ,
687
+ help = "Builder optimization level for TensorRT" ,
688
+ )
683
689
arg_parser .add_argument (
684
690
"--is_trt_engine" ,
685
691
action = "store_true" ,
0 commit comments