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

Fix export test matrices to exclude nms from Classify models (#18880)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Mohammed Yasin 2025-01-26 04:26:58 +08:00 • committed by GitHub
parent 83dc1fea6e
commit de05d1b655
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 46 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

15
ultralytics/engine/exporter.py
View file

@ -75,7 +75,7 @@ from ultralytics.data.dataset import YOLODataset
from ultralytics.data.utils import check_cls_dataset, check_det_dataset
from ultralytics.nn.autobackend import check_class_names, default_class_names
from ultralytics.nn.modules import C2f, Classify, Detect, RTDETRDecoder
from ultralytics.nn.tasks import DetectionModel, SegmentationModel, WorldModel
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, WorldModel
from ultralytics.utils import (
ARM64,
DEFAULT_CFG,
@ -282,6 +282,7 @@ class Exporter:
if self.args.int8 and tflite:
assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
if self.args.nms:
assert not isinstance(model, ClassificationModel), "'nms=True' is not valid for classification models."
if getattr(model, "end2end", False):
LOGGER.warning("WARNING ⚠️ 'nms=True' is not available for end2end models. Forcing 'nms=False'.")
self.args.nms = False
@ -507,6 +508,7 @@ class Exporter:
output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
dynamic = self.args.dynamic
if dynamic:
self.model.cpu() # dynamic=True only compatible with cpu
dynamic = {"images": {0: "batch", 2: "height", 3: "width"}} # shape(1,3,640,640)
if isinstance(self.model, SegmentationModel):
dynamic["output0"] = {0: "batch", 2: "anchors"} # shape(1, 116, 8400)
@ -518,13 +520,14 @@ class Exporter:
if self.args.nms and self.model.task == "obb":
self.args.opset = opset_version # for NMSModel
# OBB error https://github.com/pytorch/pytorch/issues/110859#issuecomment-1757841865
torch.onnx.register_custom_op_symbolic("aten::lift_fresh", lambda g, x: x, opset_version)
try:
torch.onnx.register_custom_op_symbolic("aten::lift_fresh", lambda g, x: x, opset_version)
except RuntimeError: # it will fail if it's already registered
pass
check_requirements("onnxslim>=0.1.46") # Older versions has bug with OBB
torch.onnx.export(
NMSModel(self.model.cpu() if dynamic else self.model, self.args)
if self.args.nms
else self.model, # dynamic=True only compatible with cpu
NMSModel(self.model, self.args) if self.args.nms else self.model,
self.im.cpu() if dynamic else self.im,
f,
verbose=False,
@ -1570,7 +1573,7 @@ class NMSModel(torch.nn.Module):
# TFLite GatherND error if mask is empty
score *= mask
# Explicit length otherwise reshape error, hardcoded to `self.args.max_det * 5`
mask = score.topk(self.args.max_det * 5).indices
mask = score.topk(min(self.args.max_det * 5, score.shape[0])).indices
box, score, cls, extra = box[mask], score[mask], cls[mask], extra[mask]
if not self.obb:
box = xywh2xyxy(box)