Enable ONNX --half FP16 inference (#6268)

* Enable ONNX ``--half` FP16 inference

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocherandpre-commit-ci[bot] authored and GitHub committed 2022-01-11 10:13:17 -10:00
1 parent b3eaf5008b
commit f3085accd3
3 files changed
+4 -4

No files matched your search

+2 -2
View File
@@ -137,9 +137,9 @@ def run(data,
# Load model
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine
stride, pt, jit, onnx, engine = model.stride, model.pt, model.jit, model.onnx, model.engine
imgsz = check_img_size(imgsz, s=stride) # check image size
half &= (pt or jit or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
half &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16 supported on limited backends with CUDA
if pt or jit:
model.model.half() if half else model.model.float()
elif engine: