HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Fix TFLite error and OpenVINO int8 error at imgsz=32 (#18898)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Mohammed Yasin 2025-01-26 21:18:51 +08:00 • committed by GitHub
parent 55e422d336
commit e99b778cc4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 20 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

16
ultralytics/engine/exporter.py
View file

@ -1559,6 +1559,7 @@ class NMSModel(torch.nn.Module):
extra_shape = pred.shape[-1] - (4 + self.model.nc) # extras from Segment, OBB, Pose
boxes, scores, extras = pred.split([4, self.model.nc, extra_shape], dim=2)
scores, classes = scores.max(dim=-1)
self.args.max_det = min(pred.shape[1], self.args.max_det) # in case num_anchors < max_det
# (N, max_det, 4 coords + 1 class score + 1 class label + extra_shape).
out = torch.zeros(
boxes.shape[0],
@ -1596,14 +1597,25 @@ class NMSModel(torch.nn.Module):
offbox = nmsbox[:, :end] + cls_offset * multiplier
nmsbox = torch.cat((offbox, nmsbox[:, end:]), dim=-1)
nms_fn = (
partial(nms_rotated, use_triu=not (self.is_tf or (self.args.opset or 14) < 14)) if self.obb else nms
partial(
nms_rotated,
use_triu=not (
self.is_tf
or (self.args.opset or 14) < 14
or (self.args.format == "openvino" and self.args.int8) # OpenVINO int8 error with triu
),
)
if self.obb
else nms
)
keep = nms_fn(
torch.cat([nmsbox, extra], dim=-1) if self.obb else nmsbox,
score,
self.args.iou,
)[: self.args.max_det]
dets = torch.cat([box[keep], score[keep].view(-1, 1), cls[keep].view(-1, 1), extra[keep]], dim=-1)
dets = torch.cat(
[box[keep], score[keep].view(-1, 1), cls[keep].view(-1, 1).to(out.dtype), extra[keep]], dim=-1
)
# Zero-pad to max_det size to avoid reshape error
pad = (0, 0, 0, self.args.max_det - dets.shape[0])
out[i] = torch.nn.functional.pad(dets, pad)