FP16 NMS (#814)
This commit is contained in:
1 parent
b42e8a531b
commit
ea5e97fb9b
1 file changed
-2
@@ -596,8 +596,6 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False,
|
|||||||
Returns:
|
Returns:
|
||||||
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
|
||||||
"""
|
"""
|
||||||
if prediction.dtype is torch.float16:
|
|
||||||
prediction = prediction.float() # to FP32
|
|
||||||
|
|
||||||
nc = prediction[0].shape[1] - 5 # number of classes
|
nc = prediction[0].shape[1] - 5 # number of classes
|
||||||
xc = prediction[..., 4] > conf_thres # candidates
|
xc = prediction[..., 4] > conf_thres # candidates
|
||||||
|
|||||||
Reference in new issue
Block a user