Unverified 提交 ea5e97fb authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

FP16 NMS (#814)

上级 b42e8a53
...@@ -596,8 +596,6 @@ def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, ...@@ -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
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论