diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index dac2fcf9..7fa6cf40 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/ultralytics/yolo/engine/exporter.py @@ -59,7 +59,7 @@ from pathlib import Path import torch from ultralytics.nn.autobackend import check_class_names -from ultralytics.nn.modules import C2f, Detect, Segment +from ultralytics.nn.modules import C2f, Detect, RTDETRDecoder from ultralytics.nn.tasks import DetectionModel, SegmentationModel from ultralytics.yolo.cfg import get_cfg from ultralytics.yolo.utils import (DEFAULT_CFG, LINUX, LOGGER, MACOS, __version__, callbacks, colorstr, @@ -157,13 +157,13 @@ class Exporter: # Load PyTorch model self.device = select_device('cpu' if self.args.device is None else self.args.device) + + # Checks + model.names = check_class_names(model.names) if self.args.half and onnx and self.device.type == 'cpu': LOGGER.warning('WARNING ⚠️ half=True only compatible with GPU export, i.e. use device=0') self.args.half = False assert not self.args.dynamic, 'half=True not compatible with dynamic=True, i.e. use only one.' - - # Checks - model.names = check_class_names(model.names) self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2) # check image size if self.args.optimize: assert self.device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu' @@ -185,7 +185,7 @@ class Exporter: model.float() model = model.fuse() for k, m in model.named_modules(): - if isinstance(m, (Detect, Segment)): + if isinstance(m, (Detect, RTDETRDecoder)): # Segment and Pose use Detect base class m.dynamic = self.args.dynamic m.export = True m.format = self.args.format