Default FP16 TensorRT export (#6798)

* Assert engine precision #6777

* Default to FP32 inputs for TensorRT engines

* Default to FP16 TensorRT exports #6777

* Remove wrong line #6777

* Automatically adjust detect.py input precision #6777

* Automatically adjust val.py input precision #6777

* Add missing colon

* Cleanup

* Cleanup

* Remove default trt_fp16_input definition

* Experiment

* Reorder detect.py if statement to after half checks

* Update common.py

* Update export.py

* Cleanup

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
DavidBandGlenn Jocher authored and GitHub committed 2022-03-06 21:21:16 +01:00
1 parent 7e98b4801a
commit 596de6d5a0
4 files changed
+13 -3

No files matched your search

+2 -3
View File
@@ -233,9 +233,8 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
for out in outputs:
LOGGER.info(f'{prefix}\toutput "{out.name}" with shape {out.shape} and dtype {out.dtype}')
half &= builder.platform_has_fast_fp16
LOGGER.info(f'{prefix} building FP{16 if half else 32} engine in {f}')
if half:
LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 else 32} engine in {f}')
if builder.platform_has_fast_fp16:
config.set_flag(trt.BuilderFlag.FP16)
with builder.build_engine(network, config) as engine, open(f, 'wb') as t:
t.write(engine.serialize())