Skip to content

Commit

Permalink
Automatically adjust detect.py input precision ultralytics#6777
Browse files Browse the repository at this point in the history
  • Loading branch information
DavidBaldsiefen authored Mar 5, 2022
1 parent ea07f73 commit a2c7b67
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions detect.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,9 +94,10 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
imgsz = check_img_size(imgsz, s=stride) # check image size

# Half
if engine and model.trt_fp16_input != half:
LOGGER.info('model ' + ('requires' if model.trt_fp16_input else 'incompatible with') + ' --half. Adjusting automatically.')
half = model.trt_fp16_input
half &= (pt or jit or onnx or engine) and device.type != 'cpu' # FP16 supported on limited backends with CUDA
if engine:
assert (model.trt_fp16_input == half), 'model ' + ('requires' if model.trt_fp16_input else 'incompatible with') + ' --half'
if pt or jit:
model.model.half() if half else model.model.float()

Expand Down

0 comments on commit a2c7b67

Please sign in to comment.