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:
parent
b3eaf5008b
commit
f3085accd3
|
|
@ -94,7 +94,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
||||||
imgsz = check_img_size(imgsz, s=stride) # check image size
|
imgsz = check_img_size(imgsz, s=stride) # check image size
|
||||||
|
|
||||||
# Half
|
# Half
|
||||||
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:
|
if pt or jit:
|
||||||
model.model.half() if half else model.model.float()
|
model.model.half() if half else model.model.float()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1099,4 +1099,4 @@
|
||||||
"outputs": []
|
"outputs": []
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|
|
||||||
4
val.py
4
val.py
|
|
@ -137,9 +137,9 @@ def run(data,
|
||||||
|
|
||||||
# Load model
|
# Load model
|
||||||
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
|
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
|
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:
|
if pt or jit:
|
||||||
model.model.half() if half else model.model.float()
|
model.model.half() if half else model.model.float()
|
||||||
elif engine:
|
elif engine:
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue