diff --git a/ultralytics/cfg/default.yaml b/ultralytics/cfg/default.yaml index 7b67442b..ee0f68bd 100644 --- a/ultralytics/cfg/default.yaml +++ b/ultralytics/cfg/default.yaml @@ -7,7 +7,7 @@ task: detect # (str) YOLO task, i.e. detect, segment, classify, pose, obb mode: train # (str) YOLO mode, i.e. train, val, predict, export, track, benchmark # Train settings ------------------------------------------------------------------------------------------------------- -model: # (str, optional) path to model file, i.e. yolov8n.pt, yolov8n.yaml +model: yolov13/yolov13n.pt # (str, optional) path to model file, i.e. yolov8n.pt, yolov8n.yaml data: # (str, optional) path to data file, i.e. coco8.yaml epochs: 100 # (int) number of epochs to train for time: # (float, optional) number of hours to train for, overrides epochs if supplied @@ -78,14 +78,15 @@ show_boxes: True # (bool) show prediction boxes line_width: # (int, optional) line width of the bounding boxes. Scaled to image size if None. # Export settings ------------------------------------------------------------------------------------------------------ -format: torchscript # (str) format to export to, choices at https://docs.ultralytics.com/modes/export/#export-formats +format: rknn # (str) format to export to, choices at https://docs.ultralytics.com/modes/export/#export-formats keras: False # (bool) use Kera=s optimize: False # (bool) TorchScript: optimize for mobile int8: False # (bool) CoreML/TF INT8 quantization dynamic: False # (bool) ONNX/TF/TensorRT: dynamic axes -simplify: True # (bool) ONNX: simplify model using `onnxslim` +simplify: False # (bool) ONNX: simplify model using `onnxslim` opset: # (int, optional) ONNX: opset version -workspace: None # (float, optional) TensorRT: workspace size (GiB), `None` will let TensorRT auto-allocate memory +# workspace: None # (float, optional) TensorRT: workspace size (GiB), `None` will let TensorRT auto-allocate memory +workspace: 4 # (float, optional) TensorRT: workspace size (GiB), `None` will let TensorRT auto-allocate memory nms: False # (bool) CoreML: add NMS # Hyperparameters ------------------------------------------------------------------------------------------------------ diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 37c994ef..9fa0cb25 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -53,7 +53,6 @@ $ ln -s ../../yolo11n_web_model public/yolo11n_web_model $ npm start """ - import gc import json import os @@ -122,6 +121,7 @@ def export_formats(): ["MNN", "mnn", ".mnn", True, True, ["batch", "half", "int8"]], ["NCNN", "ncnn", "_ncnn_model", True, True, ["batch", "half"]], ["IMX", "imx", "_imx_model", True, True, ["int8"]], + ['RKNN', 'rknn', '_rknnopt.torchscript', True, False, ["batch"]], ] return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x))) @@ -241,6 +241,7 @@ def __call__(self, model=None) -> str: mnn, ncnn, imx, + rknn, ) = flags # export booleans is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs)) @@ -417,6 +418,8 @@ def __call__(self, model=None) -> str: f[12], _ = self.export_ncnn() if imx: f[13], _ = self.export_imx() + if rknn: + f[12], _ = self.export_rknn() # Finish f = [str(x) for x in f if x] # filter out '' and None @@ -483,6 +486,26 @@ def export_torchscript(self, prefix=colorstr("TorchScript:")): ts.save(str(f), _extra_files=extra_files) return f, None + @try_export + def export_rknn(self, prefix=colorstr('RKNN:')): + """YOLOv13 RKNN model export.""" + LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...') + + f = str(self.file).replace(self.file.suffix, f'.onnx') + opset_version = self.args.opset or get_latest_opset() + torch.onnx.export( + self.model, + self.im[0:1,:,:,:], + f, + verbose=False, + opset_version=12, + do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False + input_names=['images']) + + LOGGER.info(f'\n{prefix} feed {f} to RKNN-Toolkit or RKNN-Toolkit2 to generate RKNN model.\n' + 'Refer https://github.com/airockchip/rknn_model_zoo/tree/main/models/CV/object_detection/yolo') + return f, None + @try_export def export_onnx(self, prefix=colorstr("ONNX:")): """YOLO ONNX export.""" @@ -1474,3 +1497,18 @@ def forward(self, x): """Normalize predictions of object detection model with input size-dependent factors.""" xywh, cls = self.model(x)[0].transpose(0, 1).split((4, self.nc), 1) return cls, xywh * self.normalize # confidence (3780, 80), coordinates (3780, 4) + +def export(cfg=DEFAULT_CFG): + """Export a YOLOv model to a specific format.""" + cfg.model = cfg.model or 'yolov8n.yaml' + cfg.format = cfg.format or 'torchscript' + from ultralytics import YOLO + model = YOLO(cfg.model) + model.export(**vars(cfg)) + +if __name__ == '__main__': + """ + CLI: + yolo mode=export model=yolov8n.yaml format=onnx + """ + export() \ No newline at end of file diff --git a/ultralytics/nn/modules/head.py b/ultralytics/nn/modules/head.py index b9d05024..c54e6db1 100644 --- a/ultralytics/nn/modules/head.py +++ b/ultralytics/nn/modules/head.py @@ -63,6 +63,16 @@ def __init__(self, nc=80, ch=()): def forward(self, x): """Concatenates and returns predicted bounding boxes and class probabilities.""" + if self.export and self.format == 'rknn': + y = [] + for i in range(self.nl): + y.append(self.cv2[i](x[i])) + clss = torch.sigmoid(self.cv3[i](x[i])) + cls_sum = torch.clamp(clss.sum(1, keepdim=True), 0, 1) + y.append(clss) + y.append(cls_sum) + return y + if self.end2end: return self.forward_end2end(x)