ultralytics 8.1.29__py3-none-any.whl → 8.3.63__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (247) hide show
  1. tests/__init__.py +22 -0
  2. tests/conftest.py +83 -0
  3. tests/test_cli.py +122 -0
  4. tests/test_cuda.py +155 -0
  5. tests/test_engine.py +131 -0
  6. tests/test_exports.py +216 -0
  7. tests/test_integrations.py +150 -0
  8. tests/test_python.py +615 -0
  9. tests/test_solutions.py +94 -0
  10. ultralytics/__init__.py +11 -8
  11. ultralytics/cfg/__init__.py +569 -131
  12. ultralytics/cfg/datasets/Argoverse.yaml +2 -1
  13. ultralytics/cfg/datasets/DOTAv1.5.yaml +3 -2
  14. ultralytics/cfg/datasets/DOTAv1.yaml +3 -2
  15. ultralytics/cfg/datasets/GlobalWheat2020.yaml +3 -2
  16. ultralytics/cfg/datasets/ImageNet.yaml +2 -1
  17. ultralytics/cfg/datasets/Objects365.yaml +5 -4
  18. ultralytics/cfg/datasets/SKU-110K.yaml +2 -1
  19. ultralytics/cfg/datasets/VOC.yaml +3 -2
  20. ultralytics/cfg/datasets/VisDrone.yaml +6 -5
  21. ultralytics/cfg/datasets/african-wildlife.yaml +25 -0
  22. ultralytics/cfg/datasets/brain-tumor.yaml +23 -0
  23. ultralytics/cfg/datasets/carparts-seg.yaml +3 -2
  24. ultralytics/cfg/datasets/coco-pose.yaml +7 -6
  25. ultralytics/cfg/datasets/coco.yaml +3 -2
  26. ultralytics/cfg/datasets/coco128-seg.yaml +4 -3
  27. ultralytics/cfg/datasets/coco128.yaml +4 -3
  28. ultralytics/cfg/datasets/coco8-pose.yaml +3 -2
  29. ultralytics/cfg/datasets/coco8-seg.yaml +3 -2
  30. ultralytics/cfg/datasets/coco8.yaml +3 -2
  31. ultralytics/cfg/datasets/crack-seg.yaml +3 -2
  32. ultralytics/cfg/datasets/dog-pose.yaml +24 -0
  33. ultralytics/cfg/datasets/dota8.yaml +3 -2
  34. ultralytics/cfg/datasets/hand-keypoints.yaml +26 -0
  35. ultralytics/cfg/datasets/lvis.yaml +1236 -0
  36. ultralytics/cfg/datasets/medical-pills.yaml +22 -0
  37. ultralytics/cfg/datasets/open-images-v7.yaml +2 -1
  38. ultralytics/cfg/datasets/package-seg.yaml +5 -4
  39. ultralytics/cfg/datasets/signature.yaml +21 -0
  40. ultralytics/cfg/datasets/tiger-pose.yaml +3 -2
  41. ultralytics/cfg/datasets/xView.yaml +2 -1
  42. ultralytics/cfg/default.yaml +14 -11
  43. ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +24 -0
  44. ultralytics/cfg/models/11/yolo11-cls.yaml +33 -0
  45. ultralytics/cfg/models/11/yolo11-obb.yaml +50 -0
  46. ultralytics/cfg/models/11/yolo11-pose.yaml +51 -0
  47. ultralytics/cfg/models/11/yolo11-seg.yaml +50 -0
  48. ultralytics/cfg/models/11/yolo11.yaml +50 -0
  49. ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +5 -2
  50. ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +5 -2
  51. ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +5 -2
  52. ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +5 -2
  53. ultralytics/cfg/models/v10/yolov10b.yaml +45 -0
  54. ultralytics/cfg/models/v10/yolov10l.yaml +45 -0
  55. ultralytics/cfg/models/v10/yolov10m.yaml +45 -0
  56. ultralytics/cfg/models/v10/yolov10n.yaml +45 -0
  57. ultralytics/cfg/models/v10/yolov10s.yaml +45 -0
  58. ultralytics/cfg/models/v10/yolov10x.yaml +45 -0
  59. ultralytics/cfg/models/v3/yolov3-spp.yaml +5 -2
  60. ultralytics/cfg/models/v3/yolov3-tiny.yaml +5 -2
  61. ultralytics/cfg/models/v3/yolov3.yaml +5 -2
  62. ultralytics/cfg/models/v5/yolov5-p6.yaml +5 -2
  63. ultralytics/cfg/models/v5/yolov5.yaml +5 -2
  64. ultralytics/cfg/models/v6/yolov6.yaml +5 -2
  65. ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +5 -2
  66. ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +5 -2
  67. ultralytics/cfg/models/v8/yolov8-cls.yaml +5 -2
  68. ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +6 -2
  69. ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +6 -2
  70. ultralytics/cfg/models/v8/yolov8-ghost.yaml +5 -2
  71. ultralytics/cfg/models/v8/yolov8-obb.yaml +5 -2
  72. ultralytics/cfg/models/v8/yolov8-p2.yaml +5 -2
  73. ultralytics/cfg/models/v8/yolov8-p6.yaml +10 -7
  74. ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +5 -2
  75. ultralytics/cfg/models/v8/yolov8-pose.yaml +5 -2
  76. ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +5 -2
  77. ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +5 -2
  78. ultralytics/cfg/models/v8/yolov8-seg.yaml +5 -2
  79. ultralytics/cfg/models/v8/yolov8-world.yaml +5 -2
  80. ultralytics/cfg/models/v8/yolov8-worldv2.yaml +5 -2
  81. ultralytics/cfg/models/v8/yolov8.yaml +5 -2
  82. ultralytics/cfg/models/v9/yolov9c-seg.yaml +41 -0
  83. ultralytics/cfg/models/v9/yolov9c.yaml +30 -25
  84. ultralytics/cfg/models/v9/yolov9e-seg.yaml +64 -0
  85. ultralytics/cfg/models/v9/yolov9e.yaml +46 -42
  86. ultralytics/cfg/models/v9/yolov9m.yaml +41 -0
  87. ultralytics/cfg/models/v9/yolov9s.yaml +41 -0
  88. ultralytics/cfg/models/v9/yolov9t.yaml +41 -0
  89. ultralytics/cfg/solutions/default.yaml +24 -0
  90. ultralytics/cfg/trackers/botsort.yaml +8 -5
  91. ultralytics/cfg/trackers/bytetrack.yaml +8 -5
  92. ultralytics/data/__init__.py +14 -3
  93. ultralytics/data/annotator.py +37 -15
  94. ultralytics/data/augment.py +1783 -289
  95. ultralytics/data/base.py +62 -27
  96. ultralytics/data/build.py +37 -8
  97. ultralytics/data/converter.py +196 -36
  98. ultralytics/data/dataset.py +233 -94
  99. ultralytics/data/loaders.py +199 -96
  100. ultralytics/data/split_dota.py +39 -29
  101. ultralytics/data/utils.py +111 -41
  102. ultralytics/engine/__init__.py +1 -1
  103. ultralytics/engine/exporter.py +579 -244
  104. ultralytics/engine/model.py +604 -252
  105. ultralytics/engine/predictor.py +22 -11
  106. ultralytics/engine/results.py +1228 -218
  107. ultralytics/engine/trainer.py +191 -129
  108. ultralytics/engine/tuner.py +18 -18
  109. ultralytics/engine/validator.py +18 -15
  110. ultralytics/hub/__init__.py +31 -13
  111. ultralytics/hub/auth.py +11 -7
  112. ultralytics/hub/google/__init__.py +159 -0
  113. ultralytics/hub/session.py +128 -94
  114. ultralytics/hub/utils.py +20 -21
  115. ultralytics/models/__init__.py +4 -2
  116. ultralytics/models/fastsam/__init__.py +2 -3
  117. ultralytics/models/fastsam/model.py +26 -4
  118. ultralytics/models/fastsam/predict.py +127 -63
  119. ultralytics/models/fastsam/utils.py +1 -44
  120. ultralytics/models/fastsam/val.py +1 -1
  121. ultralytics/models/nas/__init__.py +1 -1
  122. ultralytics/models/nas/model.py +21 -10
  123. ultralytics/models/nas/predict.py +3 -6
  124. ultralytics/models/nas/val.py +4 -4
  125. ultralytics/models/rtdetr/__init__.py +1 -1
  126. ultralytics/models/rtdetr/model.py +1 -1
  127. ultralytics/models/rtdetr/predict.py +6 -8
  128. ultralytics/models/rtdetr/train.py +6 -2
  129. ultralytics/models/rtdetr/val.py +3 -3
  130. ultralytics/models/sam/__init__.py +3 -3
  131. ultralytics/models/sam/amg.py +29 -23
  132. ultralytics/models/sam/build.py +211 -13
  133. ultralytics/models/sam/model.py +91 -30
  134. ultralytics/models/sam/modules/__init__.py +1 -1
  135. ultralytics/models/sam/modules/blocks.py +1129 -0
  136. ultralytics/models/sam/modules/decoders.py +381 -53
  137. ultralytics/models/sam/modules/encoders.py +515 -324
  138. ultralytics/models/sam/modules/memory_attention.py +237 -0
  139. ultralytics/models/sam/modules/sam.py +969 -21
  140. ultralytics/models/sam/modules/tiny_encoder.py +425 -154
  141. ultralytics/models/sam/modules/transformer.py +159 -60
  142. ultralytics/models/sam/modules/utils.py +293 -0
  143. ultralytics/models/sam/predict.py +1263 -132
  144. ultralytics/models/utils/__init__.py +1 -1
  145. ultralytics/models/utils/loss.py +36 -24
  146. ultralytics/models/utils/ops.py +3 -7
  147. ultralytics/models/yolo/__init__.py +3 -3
  148. ultralytics/models/yolo/classify/__init__.py +1 -1
  149. ultralytics/models/yolo/classify/predict.py +7 -8
  150. ultralytics/models/yolo/classify/train.py +17 -22
  151. ultralytics/models/yolo/classify/val.py +8 -4
  152. ultralytics/models/yolo/detect/__init__.py +1 -1
  153. ultralytics/models/yolo/detect/predict.py +3 -5
  154. ultralytics/models/yolo/detect/train.py +11 -4
  155. ultralytics/models/yolo/detect/val.py +90 -52
  156. ultralytics/models/yolo/model.py +14 -9
  157. ultralytics/models/yolo/obb/__init__.py +1 -1
  158. ultralytics/models/yolo/obb/predict.py +2 -2
  159. ultralytics/models/yolo/obb/train.py +5 -3
  160. ultralytics/models/yolo/obb/val.py +41 -23
  161. ultralytics/models/yolo/pose/__init__.py +1 -1
  162. ultralytics/models/yolo/pose/predict.py +3 -5
  163. ultralytics/models/yolo/pose/train.py +2 -2
  164. ultralytics/models/yolo/pose/val.py +51 -17
  165. ultralytics/models/yolo/segment/__init__.py +1 -1
  166. ultralytics/models/yolo/segment/predict.py +3 -5
  167. ultralytics/models/yolo/segment/train.py +2 -2
  168. ultralytics/models/yolo/segment/val.py +60 -19
  169. ultralytics/models/yolo/world/__init__.py +5 -0
  170. ultralytics/models/yolo/world/train.py +92 -0
  171. ultralytics/models/yolo/world/train_world.py +109 -0
  172. ultralytics/nn/__init__.py +1 -1
  173. ultralytics/nn/autobackend.py +228 -93
  174. ultralytics/nn/modules/__init__.py +39 -14
  175. ultralytics/nn/modules/activation.py +21 -0
  176. ultralytics/nn/modules/block.py +526 -66
  177. ultralytics/nn/modules/conv.py +24 -7
  178. ultralytics/nn/modules/head.py +177 -34
  179. ultralytics/nn/modules/transformer.py +6 -5
  180. ultralytics/nn/modules/utils.py +1 -2
  181. ultralytics/nn/tasks.py +226 -82
  182. ultralytics/solutions/__init__.py +30 -1
  183. ultralytics/solutions/ai_gym.py +96 -143
  184. ultralytics/solutions/analytics.py +247 -0
  185. ultralytics/solutions/distance_calculation.py +78 -135
  186. ultralytics/solutions/heatmap.py +93 -247
  187. ultralytics/solutions/object_counter.py +184 -259
  188. ultralytics/solutions/parking_management.py +246 -0
  189. ultralytics/solutions/queue_management.py +112 -0
  190. ultralytics/solutions/region_counter.py +116 -0
  191. ultralytics/solutions/security_alarm.py +144 -0
  192. ultralytics/solutions/solutions.py +178 -0
  193. ultralytics/solutions/speed_estimation.py +86 -174
  194. ultralytics/solutions/streamlit_inference.py +190 -0
  195. ultralytics/solutions/trackzone.py +68 -0
  196. ultralytics/trackers/__init__.py +1 -1
  197. ultralytics/trackers/basetrack.py +32 -13
  198. ultralytics/trackers/bot_sort.py +61 -28
  199. ultralytics/trackers/byte_tracker.py +83 -51
  200. ultralytics/trackers/track.py +21 -6
  201. ultralytics/trackers/utils/__init__.py +1 -1
  202. ultralytics/trackers/utils/gmc.py +62 -48
  203. ultralytics/trackers/utils/kalman_filter.py +166 -35
  204. ultralytics/trackers/utils/matching.py +40 -21
  205. ultralytics/utils/__init__.py +511 -239
  206. ultralytics/utils/autobatch.py +40 -22
  207. ultralytics/utils/benchmarks.py +266 -85
  208. ultralytics/utils/callbacks/__init__.py +1 -1
  209. ultralytics/utils/callbacks/base.py +1 -3
  210. ultralytics/utils/callbacks/clearml.py +7 -6
  211. ultralytics/utils/callbacks/comet.py +39 -17
  212. ultralytics/utils/callbacks/dvc.py +1 -1
  213. ultralytics/utils/callbacks/hub.py +16 -16
  214. ultralytics/utils/callbacks/mlflow.py +28 -24
  215. ultralytics/utils/callbacks/neptune.py +6 -2
  216. ultralytics/utils/callbacks/raytune.py +3 -4
  217. ultralytics/utils/callbacks/tensorboard.py +18 -18
  218. ultralytics/utils/callbacks/wb.py +27 -20
  219. ultralytics/utils/checks.py +172 -100
  220. ultralytics/utils/dist.py +2 -1
  221. ultralytics/utils/downloads.py +40 -34
  222. ultralytics/utils/errors.py +1 -1
  223. ultralytics/utils/files.py +72 -38
  224. ultralytics/utils/instance.py +41 -19
  225. ultralytics/utils/loss.py +83 -55
  226. ultralytics/utils/metrics.py +61 -56
  227. ultralytics/utils/ops.py +94 -89
  228. ultralytics/utils/patches.py +30 -14
  229. ultralytics/utils/plotting.py +600 -269
  230. ultralytics/utils/tal.py +67 -26
  231. ultralytics/utils/torch_utils.py +305 -112
  232. ultralytics/utils/triton.py +2 -1
  233. ultralytics/utils/tuner.py +21 -12
  234. ultralytics-8.3.63.dist-info/METADATA +370 -0
  235. ultralytics-8.3.63.dist-info/RECORD +241 -0
  236. {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/WHEEL +1 -1
  237. ultralytics/data/explorer/__init__.py +0 -5
  238. ultralytics/data/explorer/explorer.py +0 -472
  239. ultralytics/data/explorer/gui/__init__.py +0 -1
  240. ultralytics/data/explorer/gui/dash.py +0 -268
  241. ultralytics/data/explorer/utils.py +0 -166
  242. ultralytics/models/fastsam/prompt.py +0 -357
  243. ultralytics-8.1.29.dist-info/METADATA +0 -373
  244. ultralytics-8.1.29.dist-info/RECORD +0 -197
  245. {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/LICENSE +0 -0
  246. {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/entry_points.txt +0 -0
  247. {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/top_level.txt +0 -0
@@ -1,55 +1,60 @@
1
- # Ultralytics YOLO 🚀, AGPL-3.0 license
1
+ # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
2
  """
3
- Export a YOLOv8 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit
3
+ Export a YOLO PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit.
4
4
 
5
5
  Format | `format=argument` | Model
6
6
  --- | --- | ---
7
- PyTorch | - | yolov8n.pt
8
- TorchScript | `torchscript` | yolov8n.torchscript
9
- ONNX | `onnx` | yolov8n.onnx
10
- OpenVINO | `openvino` | yolov8n_openvino_model/
11
- TensorRT | `engine` | yolov8n.engine
12
- CoreML | `coreml` | yolov8n.mlpackage
13
- TensorFlow SavedModel | `saved_model` | yolov8n_saved_model/
14
- TensorFlow GraphDef | `pb` | yolov8n.pb
15
- TensorFlow Lite | `tflite` | yolov8n.tflite
16
- TensorFlow Edge TPU | `edgetpu` | yolov8n_edgetpu.tflite
17
- TensorFlow.js | `tfjs` | yolov8n_web_model/
18
- PaddlePaddle | `paddle` | yolov8n_paddle_model/
19
- NCNN | `ncnn` | yolov8n_ncnn_model/
7
+ PyTorch | - | yolo11n.pt
8
+ TorchScript | `torchscript` | yolo11n.torchscript
9
+ ONNX | `onnx` | yolo11n.onnx
10
+ OpenVINO | `openvino` | yolo11n_openvino_model/
11
+ TensorRT | `engine` | yolo11n.engine
12
+ CoreML | `coreml` | yolo11n.mlpackage
13
+ TensorFlow SavedModel | `saved_model` | yolo11n_saved_model/
14
+ TensorFlow GraphDef | `pb` | yolo11n.pb
15
+ TensorFlow Lite | `tflite` | yolo11n.tflite
16
+ TensorFlow Edge TPU | `edgetpu` | yolo11n_edgetpu.tflite
17
+ TensorFlow.js | `tfjs` | yolo11n_web_model/
18
+ PaddlePaddle | `paddle` | yolo11n_paddle_model/
19
+ MNN | `mnn` | yolo11n.mnn
20
+ NCNN | `ncnn` | yolo11n_ncnn_model/
21
+ IMX | `imx` | yolo11n_imx_model/
20
22
 
21
23
  Requirements:
22
24
  $ pip install "ultralytics[export]"
23
25
 
24
26
  Python:
25
27
  from ultralytics import YOLO
26
- model = YOLO('yolov8n.pt')
28
+ model = YOLO('yolo11n.pt')
27
29
  results = model.export(format='onnx')
28
30
 
29
31
  CLI:
30
- $ yolo mode=export model=yolov8n.pt format=onnx
32
+ $ yolo mode=export model=yolo11n.pt format=onnx
31
33
 
32
34
  Inference:
33
- $ yolo predict model=yolov8n.pt # PyTorch
34
- yolov8n.torchscript # TorchScript
35
- yolov8n.onnx # ONNX Runtime or OpenCV DNN with dnn=True
36
- yolov8n_openvino_model # OpenVINO
37
- yolov8n.engine # TensorRT
38
- yolov8n.mlpackage # CoreML (macOS-only)
39
- yolov8n_saved_model # TensorFlow SavedModel
40
- yolov8n.pb # TensorFlow GraphDef
41
- yolov8n.tflite # TensorFlow Lite
42
- yolov8n_edgetpu.tflite # TensorFlow Edge TPU
43
- yolov8n_paddle_model # PaddlePaddle
44
- yolov8n_ncnn_model # NCNN
35
+ $ yolo predict model=yolo11n.pt # PyTorch
36
+ yolo11n.torchscript # TorchScript
37
+ yolo11n.onnx # ONNX Runtime or OpenCV DNN with dnn=True
38
+ yolo11n_openvino_model # OpenVINO
39
+ yolo11n.engine # TensorRT
40
+ yolo11n.mlpackage # CoreML (macOS-only)
41
+ yolo11n_saved_model # TensorFlow SavedModel
42
+ yolo11n.pb # TensorFlow GraphDef
43
+ yolo11n.tflite # TensorFlow Lite
44
+ yolo11n_edgetpu.tflite # TensorFlow Edge TPU
45
+ yolo11n_paddle_model # PaddlePaddle
46
+ yolo11n.mnn # MNN
47
+ yolo11n_ncnn_model # NCNN
48
+ yolo11n_imx_model # IMX
45
49
 
46
50
  TensorFlow.js:
47
51
  $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
48
52
  $ npm install
49
- $ ln -s ../../yolov5/yolov8n_web_model public/yolov8n_web_model
53
+ $ ln -s ../../yolo11n_web_model public/yolo11n_web_model
50
54
  $ npm start
51
55
  """
52
56
 
57
+ import gc
53
58
  import json
54
59
  import os
55
60
  import shutil
@@ -63,18 +68,21 @@ from pathlib import Path
63
68
  import numpy as np
64
69
  import torch
65
70
 
66
- from ultralytics.cfg import get_cfg
71
+ from ultralytics.cfg import TASK2DATA, get_cfg
72
+ from ultralytics.data import build_dataloader
67
73
  from ultralytics.data.dataset import YOLODataset
68
- from ultralytics.data.utils import check_det_dataset
74
+ from ultralytics.data.utils import check_cls_dataset, check_det_dataset
69
75
  from ultralytics.nn.autobackend import check_class_names, default_class_names
70
- from ultralytics.nn.modules import C2f, Detect, RTDETRDecoder
76
+ from ultralytics.nn.modules import C2f, Classify, Detect, RTDETRDecoder
71
77
  from ultralytics.nn.tasks import DetectionModel, SegmentationModel, WorldModel
72
78
  from ultralytics.utils import (
73
79
  ARM64,
74
80
  DEFAULT_CFG,
81
+ IS_JETSON,
75
82
  LINUX,
76
83
  LOGGER,
77
84
  MACOS,
85
+ PYTHON_VERSION,
78
86
  ROOT,
79
87
  WINDOWS,
80
88
  __version__,
@@ -83,33 +91,63 @@ from ultralytics.utils import (
83
91
  get_default_args,
84
92
  yaml_save,
85
93
  )
86
- from ultralytics.utils.checks import PYTHON_VERSION, check_imgsz, check_is_path_safe, check_requirements, check_version
87
- from ultralytics.utils.downloads import attempt_download_asset, get_github_assets
94
+ from ultralytics.utils.checks import (
95
+ check_imgsz,
96
+ check_is_path_safe,
97
+ check_requirements,
98
+ check_version,
99
+ is_sudo_available,
100
+ )
101
+ from ultralytics.utils.downloads import attempt_download_asset, get_github_assets, safe_download
88
102
  from ultralytics.utils.files import file_size, spaces_in_path
89
103
  from ultralytics.utils.ops import Profile
90
- from ultralytics.utils.torch_utils import TORCH_1_13, get_latest_opset, select_device, smart_inference_mode
104
+ from ultralytics.utils.torch_utils import TORCH_1_13, get_latest_opset, select_device
91
105
 
92
106
 
93
107
  def export_formats():
94
- """YOLOv8 export formats."""
95
- import pandas
96
-
108
+ """Ultralytics YOLO export formats."""
97
109
  x = [
98
- ["PyTorch", "-", ".pt", True, True],
99
- ["TorchScript", "torchscript", ".torchscript", True, True],
100
- ["ONNX", "onnx", ".onnx", True, True],
101
- ["OpenVINO", "openvino", "_openvino_model", True, False],
102
- ["TensorRT", "engine", ".engine", False, True],
103
- ["CoreML", "coreml", ".mlpackage", True, False],
104
- ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True],
105
- ["TensorFlow GraphDef", "pb", ".pb", True, True],
106
- ["TensorFlow Lite", "tflite", ".tflite", True, False],
107
- ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", True, False],
108
- ["TensorFlow.js", "tfjs", "_web_model", True, False],
109
- ["PaddlePaddle", "paddle", "_paddle_model", True, True],
110
- ["NCNN", "ncnn", "_ncnn_model", True, True],
110
+ ["PyTorch", "-", ".pt", True, True, []],
111
+ ["TorchScript", "torchscript", ".torchscript", True, True, ["batch", "optimize"]],
112
+ ["ONNX", "onnx", ".onnx", True, True, ["batch", "dynamic", "half", "opset", "simplify"]],
113
+ ["OpenVINO", "openvino", "_openvino_model", True, False, ["batch", "dynamic", "half", "int8"]],
114
+ ["TensorRT", "engine", ".engine", False, True, ["batch", "dynamic", "half", "int8", "simplify"]],
115
+ ["CoreML", "coreml", ".mlpackage", True, False, ["batch", "half", "int8", "nms"]],
116
+ ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True, ["batch", "int8", "keras"]],
117
+ ["TensorFlow GraphDef", "pb", ".pb", True, True, ["batch"]],
118
+ ["TensorFlow Lite", "tflite", ".tflite", True, False, ["batch", "half", "int8"]],
119
+ ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", True, False, []],
120
+ ["TensorFlow.js", "tfjs", "_web_model", True, False, ["batch", "half", "int8"]],
121
+ ["PaddlePaddle", "paddle", "_paddle_model", True, True, ["batch"]],
122
+ ["MNN", "mnn", ".mnn", True, True, ["batch", "half", "int8"]],
123
+ ["NCNN", "ncnn", "_ncnn_model", True, True, ["batch", "half"]],
124
+ ["IMX", "imx", "_imx_model", True, True, ["int8"]],
111
125
  ]
112
- return pandas.DataFrame(x, columns=["Format", "Argument", "Suffix", "CPU", "GPU"])
126
+ return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
127
+
128
+
129
+ def validate_args(format, passed_args, valid_args):
130
+ """
131
+ Validates arguments based on format.
132
+
133
+ Args:
134
+ format (str): The export format.
135
+ passed_args (Namespace): The arguments used during export.
136
+ valid_args (dict): List of valid arguments for the format.
137
+
138
+ Raises:
139
+ AssertionError: If an argument that's not supported by the export format is used, or if format doesn't have the supported arguments listed.
140
+ """
141
+ # Only check valid usage of these args
142
+ export_args = ["half", "int8", "dynamic", "keras", "nms", "batch"]
143
+
144
+ assert valid_args is not None, f"ERROR ❌️ valid arguments for '{format}' not listed."
145
+ custom = {"batch": 1, "data": None, "device": None} # exporter defaults
146
+ default_args = get_cfg(DEFAULT_CFG, custom)
147
+ for arg in export_args:
148
+ not_default = getattr(passed_args, arg, None) != getattr(default_args, arg, None)
149
+ if not_default:
150
+ assert arg in valid_args, f"ERROR ❌️ argument '{arg}' is not supported for format='{format}'"
113
151
 
114
152
 
115
153
  def gd_outputs(gd):
@@ -122,7 +160,7 @@ def gd_outputs(gd):
122
160
 
123
161
 
124
162
  def try_export(inner_func):
125
- """YOLOv8 export decorator, i..e @try_export."""
163
+ """YOLO export decorator, i.e. @try_export."""
126
164
  inner_args = get_default_args(inner_func)
127
165
 
128
166
  def outer_func(*args, **kwargs):
@@ -134,7 +172,7 @@ def try_export(inner_func):
134
172
  LOGGER.info(f"{prefix} export success ✅ {dt.t:.1f}s, saved as '{f}' ({file_size(f):.1f} MB)")
135
173
  return f, model
136
174
  except Exception as e:
137
- LOGGER.info(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}")
175
+ LOGGER.error(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}")
138
176
  raise e
139
177
 
140
178
  return outer_func
@@ -159,48 +197,94 @@ class Exporter:
159
197
  _callbacks (dict, optional): Dictionary of callback functions. Defaults to None.
160
198
  """
161
199
  self.args = get_cfg(cfg, overrides)
162
- if self.args.format.lower() in ("coreml", "mlmodel"): # fix attempt for protobuf<3.20.x errors
200
+ if self.args.format.lower() in {"coreml", "mlmodel"}: # fix attempt for protobuf<3.20.x errors
163
201
  os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" # must run before TensorBoard callback
164
202
 
165
203
  self.callbacks = _callbacks or callbacks.get_default_callbacks()
166
204
  callbacks.add_integration_callbacks(self)
167
205
 
168
- @smart_inference_mode()
169
- def __call__(self, model=None):
206
+ def __call__(self, model=None) -> str:
170
207
  """Returns list of exported files/dirs after running callbacks."""
171
208
  self.run_callbacks("on_export_start")
172
209
  t = time.time()
173
210
  fmt = self.args.format.lower() # to lowercase
174
- if fmt in ("tensorrt", "trt"): # 'engine' aliases
211
+ if fmt in {"tensorrt", "trt"}: # 'engine' aliases
175
212
  fmt = "engine"
176
- if fmt in ("mlmodel", "mlpackage", "mlprogram", "apple", "ios", "coreml"): # 'coreml' aliases
213
+ if fmt in {"mlmodel", "mlpackage", "mlprogram", "apple", "ios", "coreml"}: # 'coreml' aliases
177
214
  fmt = "coreml"
178
- fmts = tuple(export_formats()["Argument"][1:]) # available export formats
215
+ fmts_dict = export_formats()
216
+ fmts = tuple(fmts_dict["Argument"][1:]) # available export formats
217
+ if fmt not in fmts:
218
+ import difflib
219
+
220
+ # Get the closest match if format is invalid
221
+ matches = difflib.get_close_matches(fmt, fmts, n=1, cutoff=0.6) # 60% similarity required to match
222
+ if not matches:
223
+ raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
224
+ LOGGER.warning(f"WARNING ⚠️ Invalid export format='{fmt}', updating to format='{matches[0]}'")
225
+ fmt = matches[0]
179
226
  flags = [x == fmt for x in fmts]
180
227
  if sum(flags) != 1:
181
228
  raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
182
- jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn = flags # export booleans
229
+ (
230
+ jit,
231
+ onnx,
232
+ xml,
233
+ engine,
234
+ coreml,
235
+ saved_model,
236
+ pb,
237
+ tflite,
238
+ edgetpu,
239
+ tfjs,
240
+ paddle,
241
+ mnn,
242
+ ncnn,
243
+ imx,
244
+ ) = flags # export booleans
245
+ is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
183
246
 
184
247
  # Device
248
+ dla = None
185
249
  if fmt == "engine" and self.args.device is None:
186
250
  LOGGER.warning("WARNING ⚠️ TensorRT requires GPU export, automatically assigning device=0")
187
251
  self.args.device = "0"
252
+ if fmt == "engine" and "dla" in str(self.args.device): # convert int/list to str first
253
+ dla = self.args.device.split(":")[-1]
254
+ self.args.device = "0" # update device to "0"
255
+ assert dla in {"0", "1"}, f"Expected self.args.device='dla:0' or 'dla:1, but got {self.args.device}."
188
256
  self.device = select_device("cpu" if self.args.device is None else self.args.device)
189
257
 
190
- # Checks
258
+ # Argument compatibility checks
259
+ fmt_keys = fmts_dict["Arguments"][flags.index(True) + 1]
260
+ validate_args(fmt, self.args, fmt_keys)
261
+ if imx and not self.args.int8:
262
+ LOGGER.warning("WARNING ⚠️ IMX only supports int8 export, setting int8=True.")
263
+ self.args.int8 = True
191
264
  if not hasattr(model, "names"):
192
265
  model.names = default_class_names()
193
266
  model.names = check_class_names(model.names)
267
+ if self.args.half and self.args.int8:
268
+ LOGGER.warning("WARNING ⚠️ half=True and int8=True are mutually exclusive, setting half=False.")
269
+ self.args.half = False
194
270
  if self.args.half and onnx and self.device.type == "cpu":
195
271
  LOGGER.warning("WARNING ⚠️ half=True only compatible with GPU export, i.e. use device=0")
196
272
  self.args.half = False
197
273
  assert not self.args.dynamic, "half=True not compatible with dynamic=True, i.e. use only one."
198
274
  self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2) # check image size
275
+ if self.args.int8 and engine:
276
+ self.args.dynamic = True # enforce dynamic to export TensorRT INT8
199
277
  if self.args.optimize:
200
278
  assert not ncnn, "optimize=True not compatible with format='ncnn', i.e. use optimize=False"
201
279
  assert self.device.type == "cpu", "optimize=True not compatible with cuda devices, i.e. use device='cpu'"
202
- if edgetpu and not LINUX:
203
- raise SystemError("Edge TPU export only supported on Linux. See https://coral.ai/docs/edgetpu/compiler/")
280
+ if self.args.int8 and tflite:
281
+ assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
282
+ if edgetpu:
283
+ if not LINUX:
284
+ raise SystemError("Edge TPU export only supported on Linux. See https://coral.ai/docs/edgetpu/compiler")
285
+ elif self.args.batch != 1: # see github.com/ultralytics/ultralytics/pull/13420
286
+ LOGGER.warning("WARNING ⚠️ Edge TPU export requires batch size 1, setting batch=1.")
287
+ self.args.batch = 1
204
288
  if isinstance(model, WorldModel):
205
289
  LOGGER.warning(
206
290
  "WARNING ⚠️ YOLOWorld (original version) export is not supported to any format.\n"
@@ -208,6 +292,13 @@ class Exporter:
208
292
  "(torchscript, onnx, openvino, engine, coreml) formats. "
209
293
  "See https://docs.ultralytics.com/models/yolo-world for details."
210
294
  )
295
+ model.clip_model = None # openvino int8 export error: https://github.com/ultralytics/ultralytics/pull/18445
296
+ if self.args.int8 and not self.args.data:
297
+ self.args.data = DEFAULT_CFG.data or TASK2DATA[getattr(model, "task", "detect")] # assign default data
298
+ LOGGER.warning(
299
+ "WARNING ⚠️ INT8 export requires a missing 'data' arg for calibration. "
300
+ f"Using default 'data={self.args.data}'."
301
+ )
211
302
 
212
303
  # Input
213
304
  im = torch.zeros(self.args.batch, 3, *self.imgsz).to(self.device)
@@ -224,14 +315,31 @@ class Exporter:
224
315
  model.eval()
225
316
  model.float()
226
317
  model = model.fuse()
318
+
319
+ if imx:
320
+ from ultralytics.utils.torch_utils import FXModel
321
+
322
+ model = FXModel(model)
227
323
  for m in model.modules():
324
+ if isinstance(m, Classify):
325
+ m.export = True
228
326
  if isinstance(m, (Detect, RTDETRDecoder)): # includes all Detect subclasses like Segment, Pose, OBB
229
327
  m.dynamic = self.args.dynamic
230
328
  m.export = True
231
329
  m.format = self.args.format
232
- elif isinstance(m, C2f) and not any((saved_model, pb, tflite, edgetpu, tfjs)):
330
+ m.max_det = self.args.max_det
331
+ elif isinstance(m, C2f) and not is_tf_format:
233
332
  # EdgeTPU does not support FlexSplitV while split provides cleaner ONNX graph
234
333
  m.forward = m.forward_split
334
+ if isinstance(m, Detect) and imx:
335
+ from ultralytics.utils.tal import make_anchors
336
+
337
+ m.anchors, m.strides = (
338
+ x.transpose(0, 1)
339
+ for x in make_anchors(
340
+ torch.cat([s / m.stride.unsqueeze(-1) for s in self.imgsz], dim=1), m.stride, 0.5
341
+ )
342
+ )
235
343
 
236
344
  y = None
237
345
  for _ in range(2):
@@ -255,7 +363,7 @@ class Exporter:
255
363
  )
256
364
  self.pretty_name = Path(self.model.yaml.get("yaml_file", self.file)).stem.replace("yolo", "YOLO")
257
365
  data = model.args["data"] if hasattr(model, "args") and isinstance(model.args, dict) else ""
258
- description = f'Ultralytics {self.pretty_name} model {f"trained on {data}" if data else ""}'
366
+ description = f"Ultralytics {self.pretty_name} model {f'trained on {data}' if data else ''}"
259
367
  self.metadata = {
260
368
  "description": description,
261
369
  "author": "Ultralytics",
@@ -268,13 +376,14 @@ class Exporter:
268
376
  "batch": self.args.batch,
269
377
  "imgsz": self.imgsz,
270
378
  "names": model.names,
379
+ "args": {k: v for k, v in self.args if k in fmt_keys},
271
380
  } # model metadata
272
381
  if model.task == "pose":
273
382
  self.metadata["kpt_shape"] = model.model[-1].kpt_shape
274
383
 
275
384
  LOGGER.info(
276
385
  f"\n{colorstr('PyTorch:')} starting from '{file}' with input shape {tuple(im.shape)} BCHW and "
277
- f'output shape(s) {self.output_shape} ({file_size(file):.1f} MB)'
386
+ f"output shape(s) {self.output_shape} ({file_size(file):.1f} MB)"
278
387
  )
279
388
 
280
389
  # Exports
@@ -282,14 +391,14 @@ class Exporter:
282
391
  if jit or ncnn: # TorchScript
283
392
  f[0], _ = self.export_torchscript()
284
393
  if engine: # TensorRT required before ONNX
285
- f[1], _ = self.export_engine()
394
+ f[1], _ = self.export_engine(dla=dla)
286
395
  if onnx: # ONNX
287
396
  f[2], _ = self.export_onnx()
288
397
  if xml: # OpenVINO
289
398
  f[3], _ = self.export_openvino()
290
399
  if coreml: # CoreML
291
400
  f[4], _ = self.export_coreml()
292
- if any((saved_model, pb, tflite, edgetpu, tfjs)): # TensorFlow formats
401
+ if is_tf_format: # TensorFlow formats
293
402
  self.args.int8 |= edgetpu
294
403
  f[5], keras_model = self.export_saved_model()
295
404
  if pb or tfjs: # pb prerequisite to tfjs
@@ -302,8 +411,12 @@ class Exporter:
302
411
  f[9], _ = self.export_tfjs()
303
412
  if paddle: # PaddlePaddle
304
413
  f[10], _ = self.export_paddle()
414
+ if mnn: # MNN
415
+ f[11], _ = self.export_mnn()
305
416
  if ncnn: # NCNN
306
- f[11], _ = self.export_ncnn()
417
+ f[12], _ = self.export_ncnn()
418
+ if imx:
419
+ f[13], _ = self.export_imx()
307
420
 
308
421
  # Finish
309
422
  f = [str(x) for x in f if x] # filter out '' and None
@@ -320,19 +433,42 @@ class Exporter:
320
433
  predict_data = f"data={data}" if model.task == "segment" and fmt == "pb" else ""
321
434
  q = "int8" if self.args.int8 else "half" if self.args.half else "" # quantization
322
435
  LOGGER.info(
323
- f'\nExport complete ({time.time() - t:.1f}s)'
436
+ f"\nExport complete ({time.time() - t:.1f}s)"
324
437
  f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
325
- f'\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {q} {predict_data}'
326
- f'\nValidate: yolo val task={model.task} model={f} imgsz={imgsz} data={data} {q} {s}'
327
- f'\nVisualize: https://netron.app'
438
+ f"\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {q} {predict_data}"
439
+ f"\nValidate: yolo val task={model.task} model={f} imgsz={imgsz} data={data} {q} {s}"
440
+ f"\nVisualize: https://netron.app"
328
441
  )
329
442
 
330
443
  self.run_callbacks("on_export_end")
331
444
  return f # return list of exported files/dirs
332
445
 
446
+ def get_int8_calibration_dataloader(self, prefix=""):
447
+ """Build and return a dataloader suitable for calibration of INT8 models."""
448
+ LOGGER.info(f"{prefix} collecting INT8 calibration images from 'data={self.args.data}'")
449
+ data = (check_cls_dataset if self.model.task == "classify" else check_det_dataset)(self.args.data)
450
+ # TensorRT INT8 calibration should use 2x batch size
451
+ batch = self.args.batch * (2 if self.args.format == "engine" else 1)
452
+ dataset = YOLODataset(
453
+ data[self.args.split or "val"],
454
+ data=data,
455
+ task=self.model.task,
456
+ imgsz=self.imgsz[0],
457
+ augment=False,
458
+ batch_size=batch,
459
+ )
460
+ n = len(dataset)
461
+ if n < self.args.batch:
462
+ raise ValueError(
463
+ f"The calibration dataset ({n} images) must have at least as many images as the batch size ('batch={self.args.batch}')."
464
+ )
465
+ elif n < 300:
466
+ LOGGER.warning(f"{prefix} WARNING ⚠️ >300 images recommended for INT8 calibration, found {n} images.")
467
+ return build_dataloader(dataset, batch=batch, workers=0) # required for batch loading
468
+
333
469
  @try_export
334
470
  def export_torchscript(self, prefix=colorstr("TorchScript:")):
335
- """YOLOv8 TorchScript model export."""
471
+ """YOLO TorchScript model export."""
336
472
  LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
337
473
  f = self.file.with_suffix(".torchscript")
338
474
 
@@ -349,12 +485,10 @@ class Exporter:
349
485
 
350
486
  @try_export
351
487
  def export_onnx(self, prefix=colorstr("ONNX:")):
352
- """YOLOv8 ONNX export."""
488
+ """YOLO ONNX export."""
353
489
  requirements = ["onnx>=1.12.0"]
354
490
  if self.args.simplify:
355
- requirements += ["onnxsim>=0.4.33", "onnxruntime-gpu" if torch.cuda.is_available() else "onnxruntime"]
356
- if ARM64:
357
- check_requirements("cmake") # 'cmake' is needed to build onnxsim on aarch64
491
+ requirements += ["onnxslim", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
358
492
  check_requirements(requirements)
359
493
  import onnx # noqa
360
494
 
@@ -386,19 +520,17 @@ class Exporter:
386
520
 
387
521
  # Checks
388
522
  model_onnx = onnx.load(f) # load onnx model
389
- # onnx.checker.check_model(model_onnx) # check onnx model
390
523
 
391
524
  # Simplify
392
525
  if self.args.simplify:
393
526
  try:
394
- import onnxsim
527
+ import onnxslim
528
+
529
+ LOGGER.info(f"{prefix} slimming with onnxslim {onnxslim.__version__}...")
530
+ model_onnx = onnxslim.slim(model_onnx)
395
531
 
396
- LOGGER.info(f"{prefix} simplifying with onnxsim {onnxsim.__version__}...")
397
- # subprocess.run(f'onnxsim "{f}" "{f}"', shell=True)
398
- model_onnx, check = onnxsim.simplify(model_onnx)
399
- assert check, "Simplified ONNX model could not be validated"
400
532
  except Exception as e:
401
- LOGGER.info(f"{prefix} simplifier failure: {e}")
533
+ LOGGER.warning(f"{prefix} simplifier failure: {e}")
402
534
 
403
535
  # Metadata
404
536
  for k, v in self.metadata.items():
@@ -410,21 +542,21 @@ class Exporter:
410
542
 
411
543
  @try_export
412
544
  def export_openvino(self, prefix=colorstr("OpenVINO:")):
413
- """YOLOv8 OpenVINO export."""
414
- check_requirements("openvino>=2024.0.0") # requires openvino: https://pypi.org/project/openvino/
545
+ """YOLO OpenVINO export."""
546
+ check_requirements("openvino>=2024.5.0")
415
547
  import openvino as ov
416
548
 
417
549
  LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
418
550
  assert TORCH_1_13, f"OpenVINO export requires torch>=1.13.0 but torch=={torch.__version__} is installed"
419
551
  ov_model = ov.convert_model(
420
- self.model.cpu(),
552
+ self.model,
421
553
  input=None if self.args.dynamic else [self.im.shape],
422
554
  example_input=self.im,
423
555
  )
424
556
 
425
557
  def serialize(ov_model, file):
426
558
  """Set RT info, serialize and save metadata YAML."""
427
- ov_model.set_rt_info("YOLOv8", ["model_info", "model_type"])
559
+ ov_model.set_rt_info("YOLO", ["model_info", "model_type"])
428
560
  ov_model.set_rt_info(True, ["model_info", "reverse_input_channels"])
429
561
  ov_model.set_rt_info(114, ["model_info", "pad_value"])
430
562
  ov_model.set_rt_info([255.0], ["model_info", "scale_values"])
@@ -439,37 +571,21 @@ class Exporter:
439
571
  if self.args.int8:
440
572
  fq = str(self.file).replace(self.file.suffix, f"_int8_openvino_model{os.sep}")
441
573
  fq_ov = str(Path(fq) / self.file.with_suffix(".xml").name)
442
- if not self.args.data:
443
- self.args.data = DEFAULT_CFG.data or "coco128.yaml"
444
- LOGGER.warning(
445
- f"{prefix} WARNING ⚠️ INT8 export requires a missing 'data' arg for calibration. "
446
- f"Using default 'data={self.args.data}'."
447
- )
448
- check_requirements("nncf>=2.8.0")
574
+ check_requirements("nncf>=2.14.0")
449
575
  import nncf
450
576
 
451
- def transform_fn(data_item):
577
+ def transform_fn(data_item) -> np.ndarray:
452
578
  """Quantization transform function."""
453
- assert (
454
- data_item["img"].dtype == torch.uint8
455
- ), "Input image must be uint8 for the quantization preprocessing"
456
- im = data_item["img"].numpy().astype(np.float32) / 255.0 # uint8 to fp16/32 and 0 - 255 to 0.0 - 1.0
579
+ data_item: torch.Tensor = data_item["img"] if isinstance(data_item, dict) else data_item
580
+ assert data_item.dtype == torch.uint8, "Input image must be uint8 for the quantization preprocessing"
581
+ im = data_item.numpy().astype(np.float32) / 255.0 # uint8 to fp16/32 and 0 - 255 to 0.0 - 1.0
457
582
  return np.expand_dims(im, 0) if im.ndim == 3 else im
458
583
 
459
584
  # Generate calibration data for integer quantization
460
- LOGGER.info(f"{prefix} collecting INT8 calibration images from 'data={self.args.data}'")
461
- data = check_det_dataset(self.args.data)
462
- dataset = YOLODataset(data["val"], data=data, imgsz=self.imgsz[0], augment=False)
463
- n = len(dataset)
464
- if n < 300:
465
- LOGGER.warning(f"{prefix} WARNING ⚠️ >300 images recommended for INT8 calibration, found {n} images.")
466
- quantization_dataset = nncf.Dataset(dataset, transform_fn)
467
-
468
585
  ignored_scope = None
469
586
  if isinstance(self.model.model[-1], Detect):
470
587
  # Includes all Detect subclasses like Segment, Pose, OBB, WorldDetect
471
588
  head_module_name = ".".join(list(self.model.named_modules())[-1][0].split(".")[:2])
472
-
473
589
  ignored_scope = nncf.IgnoredScope( # ignore operations
474
590
  patterns=[
475
591
  f".*{head_module_name}/.*/Add",
@@ -482,7 +598,10 @@ class Exporter:
482
598
  )
483
599
 
484
600
  quantized_ov_model = nncf.quantize(
485
- ov_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED, ignored_scope=ignored_scope
601
+ model=ov_model,
602
+ calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix), transform_fn),
603
+ preset=nncf.QuantizationPreset.MIXED,
604
+ ignored_scope=ignored_scope,
486
605
  )
487
606
  serialize(quantized_ov_model, fq_ov)
488
607
  return fq, None
@@ -495,8 +614,8 @@ class Exporter:
495
614
 
496
615
  @try_export
497
616
  def export_paddle(self, prefix=colorstr("PaddlePaddle:")):
498
- """YOLOv8 Paddle export."""
499
- check_requirements(("paddlepaddle", "x2paddle"))
617
+ """YOLO Paddle export."""
618
+ check_requirements(("paddlepaddle-gpu" if torch.cuda.is_available() else "paddlepaddle", "x2paddle"))
500
619
  import x2paddle # noqa
501
620
  from x2paddle.convert import pytorch2paddle # noqa
502
621
 
@@ -507,11 +626,34 @@ class Exporter:
507
626
  yaml_save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
508
627
  return f, None
509
628
 
629
+ @try_export
630
+ def export_mnn(self, prefix=colorstr("MNN:")):
631
+ """YOLOv8 MNN export using MNN https://github.com/alibaba/MNN."""
632
+ f_onnx, _ = self.export_onnx() # get onnx model first
633
+
634
+ check_requirements("MNN>=2.9.6")
635
+ import MNN # noqa
636
+ from MNN.tools import mnnconvert
637
+
638
+ # Setup and checks
639
+ LOGGER.info(f"\n{prefix} starting export with MNN {MNN.version()}...")
640
+ assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
641
+ f = str(self.file.with_suffix(".mnn")) # MNN model file
642
+ args = ["", "-f", "ONNX", "--modelFile", f_onnx, "--MNNModel", f, "--bizCode", json.dumps(self.metadata)]
643
+ if self.args.int8:
644
+ args.extend(("--weightQuantBits", "8"))
645
+ if self.args.half:
646
+ args.append("--fp16")
647
+ mnnconvert.convert(args)
648
+ # remove scratch file for model convert optimize
649
+ convert_scratch = Path(self.file.parent / ".__convert_external_data.bin")
650
+ if convert_scratch.exists():
651
+ convert_scratch.unlink()
652
+ return f, None
653
+
510
654
  @try_export
511
655
  def export_ncnn(self, prefix=colorstr("NCNN:")):
512
- """
513
- YOLOv8 NCNN export using PNNX https://github.com/pnnx/pnnx.
514
- """
656
+ """YOLO NCNN export using PNNX https://github.com/pnnx/pnnx."""
515
657
  check_requirements("ncnn")
516
658
  import ncnn # noqa
517
659
 
@@ -520,7 +662,7 @@ class Exporter:
520
662
  f_ts = self.file.with_suffix(".torchscript")
521
663
 
522
664
  name = Path("pnnx.exe" if WINDOWS else "pnnx") # PNNX filename
523
- pnnx = name if name.is_file() else ROOT / name
665
+ pnnx = name if name.is_file() else (ROOT / name)
524
666
  if not pnnx.is_file():
525
667
  LOGGER.warning(
526
668
  f"{prefix} WARNING ⚠️ PNNX not found. Attempting to download binary file from "
@@ -528,31 +670,32 @@ class Exporter:
528
670
  f"or in {ROOT}. See PNNX repo for full installation instructions."
529
671
  )
530
672
  system = "macos" if MACOS else "windows" if WINDOWS else "linux-aarch64" if ARM64 else "linux"
531
- _, assets = get_github_assets(repo="pnnx/pnnx", retry=True)
532
- if assets:
533
- url = [x for x in assets if f"{system}.zip" in x][0]
534
- else:
535
- url = f"https://github.com/pnnx/pnnx/releases/download/20240226/pnnx-20240226-{system}.zip"
536
- LOGGER.warning(f"{prefix} WARNING ⚠️ PNNX GitHub assets not found, using default {url}")
537
- asset = attempt_download_asset(url, repo="pnnx/pnnx", release="latest")
538
- if check_is_path_safe(Path.cwd(), asset): # avoid path traversal security vulnerability
539
- unzip_dir = Path(asset).with_suffix("")
540
- (unzip_dir / name).rename(pnnx) # move binary to ROOT
541
- shutil.rmtree(unzip_dir) # delete unzip dir
542
- Path(asset).unlink() # delete zip
673
+ try:
674
+ release, assets = get_github_assets(repo="pnnx/pnnx")
675
+ asset = [x for x in assets if f"{system}.zip" in x][0]
676
+ assert isinstance(asset, str), "Unable to retrieve PNNX repo assets" # i.e. pnnx-20240410-macos.zip
677
+ LOGGER.info(f"{prefix} successfully found latest PNNX asset file {asset}")
678
+ except Exception as e:
679
+ release = "20240410"
680
+ asset = f"pnnx-{release}-{system}.zip"
681
+ LOGGER.warning(f"{prefix} WARNING ⚠️ PNNX GitHub assets not found: {e}, using default {asset}")
682
+ unzip_dir = safe_download(f"https://github.com/pnnx/pnnx/releases/download/{release}/{asset}", delete=True)
683
+ if check_is_path_safe(Path.cwd(), unzip_dir): # avoid path traversal security vulnerability
684
+ shutil.move(src=unzip_dir / name, dst=pnnx) # move binary to ROOT
543
685
  pnnx.chmod(0o777) # set read, write, and execute permissions for everyone
686
+ shutil.rmtree(unzip_dir) # delete unzip dir
544
687
 
545
688
  ncnn_args = [
546
- f'ncnnparam={f / "model.ncnn.param"}',
547
- f'ncnnbin={f / "model.ncnn.bin"}',
548
- f'ncnnpy={f / "model_ncnn.py"}',
689
+ f"ncnnparam={f / 'model.ncnn.param'}",
690
+ f"ncnnbin={f / 'model.ncnn.bin'}",
691
+ f"ncnnpy={f / 'model_ncnn.py'}",
549
692
  ]
550
693
 
551
694
  pnnx_args = [
552
- f'pnnxparam={f / "model.pnnx.param"}',
553
- f'pnnxbin={f / "model.pnnx.bin"}',
554
- f'pnnxpy={f / "model_pnnx.py"}',
555
- f'pnnxonnx={f / "model.pnnx.onnx"}',
695
+ f"pnnxparam={f / 'model.pnnx.param'}",
696
+ f"pnnxbin={f / 'model.pnnx.bin'}",
697
+ f"pnnxpy={f / 'model_pnnx.py'}",
698
+ f"pnnxonnx={f / 'model.pnnx.onnx'}",
556
699
  ]
557
700
 
558
701
  cmd = [
@@ -578,16 +721,20 @@ class Exporter:
578
721
 
579
722
  @try_export
580
723
  def export_coreml(self, prefix=colorstr("CoreML:")):
581
- """YOLOv8 CoreML export."""
724
+ """YOLO CoreML export."""
582
725
  mlmodel = self.args.format.lower() == "mlmodel" # legacy *.mlmodel export format requested
583
726
  check_requirements("coremltools>=6.0,<=6.2" if mlmodel else "coremltools>=7.0")
584
727
  import coremltools as ct # noqa
585
728
 
586
729
  LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
587
730
  assert not WINDOWS, "CoreML export is not supported on Windows, please run on macOS or Linux."
731
+ assert self.args.batch == 1, "CoreML batch sizes > 1 are not supported. Please retry at 'batch=1'."
588
732
  f = self.file.with_suffix(".mlmodel" if mlmodel else ".mlpackage")
589
733
  if f.is_dir():
590
734
  shutil.rmtree(f)
735
+ if self.args.nms and getattr(self.model, "end2end", False):
736
+ LOGGER.warning(f"{prefix} WARNING ⚠️ 'nms=True' is not available for end2end models. Forcing 'nms=False'.")
737
+ self.args.nms = False
591
738
 
592
739
  bias = [0.0, 0.0, 0.0]
593
740
  scale = 1 / 255
@@ -650,40 +797,61 @@ class Exporter:
650
797
  return f, ct_model
651
798
 
652
799
  @try_export
653
- def export_engine(self, prefix=colorstr("TensorRT:")):
654
- """YOLOv8 TensorRT export https://developer.nvidia.com/tensorrt."""
800
+ def export_engine(self, dla=None, prefix=colorstr("TensorRT:")):
801
+ """YOLO TensorRT export https://developer.nvidia.com/tensorrt."""
655
802
  assert self.im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. use 'device=0'"
656
- f_onnx, _ = self.export_onnx() # run before trt import https://github.com/ultralytics/ultralytics/issues/7016
803
+ f_onnx, _ = self.export_onnx() # run before TRT import https://github.com/ultralytics/ultralytics/issues/7016
657
804
 
658
805
  try:
659
806
  import tensorrt as trt # noqa
660
807
  except ImportError:
661
808
  if LINUX:
662
- check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com")
809
+ check_requirements("tensorrt>7.0.0,!=10.1.0")
663
810
  import tensorrt as trt # noqa
811
+ check_version(trt.__version__, ">=7.0.0", hard=True)
812
+ check_version(trt.__version__, "!=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
664
813
 
665
- check_version(trt.__version__, "7.0.0", hard=True) # require tensorrt>=7.0.0
666
-
667
- self.args.simplify = True
668
-
814
+ # Setup and checks
669
815
  LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")
816
+ is_trt10 = int(trt.__version__.split(".")[0]) >= 10 # is TensorRT >= 10
670
817
  assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
671
818
  f = self.file.with_suffix(".engine") # TensorRT engine file
672
819
  logger = trt.Logger(trt.Logger.INFO)
673
820
  if self.args.verbose:
674
821
  logger.min_severity = trt.Logger.Severity.VERBOSE
675
822
 
823
+ # Engine builder
676
824
  builder = trt.Builder(logger)
677
825
  config = builder.create_builder_config()
678
- config.max_workspace_size = self.args.workspace * 1 << 30
679
- # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice
680
-
826
+ workspace = int(self.args.workspace * (1 << 30)) if self.args.workspace is not None else 0
827
+ if is_trt10 and workspace > 0:
828
+ config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace)
829
+ elif workspace > 0: # TensorRT versions 7, 8
830
+ config.max_workspace_size = workspace
681
831
  flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
682
832
  network = builder.create_network(flag)
833
+ half = builder.platform_has_fast_fp16 and self.args.half
834
+ int8 = builder.platform_has_fast_int8 and self.args.int8
835
+
836
+ # Optionally switch to DLA if enabled
837
+ if dla is not None:
838
+ if not IS_JETSON:
839
+ raise ValueError("DLA is only available on NVIDIA Jetson devices")
840
+ LOGGER.info(f"{prefix} enabling DLA on core {dla}...")
841
+ if not self.args.half and not self.args.int8:
842
+ raise ValueError(
843
+ "DLA requires either 'half=True' (FP16) or 'int8=True' (INT8) to be enabled. Please enable one of them and try again."
844
+ )
845
+ config.default_device_type = trt.DeviceType.DLA
846
+ config.DLA_core = int(dla)
847
+ config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
848
+
849
+ # Read ONNX file
683
850
  parser = trt.OnnxParser(network, logger)
684
851
  if not parser.parse_from_file(f_onnx):
685
852
  raise RuntimeError(f"failed to load ONNX file: {f_onnx}")
686
853
 
854
+ # Network inputs
687
855
  inputs = [network.get_input(i) for i in range(network.num_inputs)]
688
856
  outputs = [network.get_output(i) for i in range(network.num_outputs)]
689
857
  for inp in inputs:
@@ -696,61 +864,117 @@ class Exporter:
696
864
  if shape[0] <= 1:
697
865
  LOGGER.warning(f"{prefix} WARNING ⚠️ 'dynamic=True' model requires max batch size, i.e. 'batch=16'")
698
866
  profile = builder.create_optimization_profile()
867
+ min_shape = (1, shape[1], 32, 32) # minimum input shape
868
+ max_shape = (*shape[:2], *(int(max(1, workspace) * d) for d in shape[2:])) # max input shape
699
869
  for inp in inputs:
700
- profile.set_shape(inp.name, (1, *shape[1:]), (max(1, shape[0] // 2), *shape[1:]), shape)
870
+ profile.set_shape(inp.name, min=min_shape, opt=shape, max=max_shape)
701
871
  config.add_optimization_profile(profile)
702
872
 
703
- LOGGER.info(
704
- f"{prefix} building FP{16 if builder.platform_has_fast_fp16 and self.args.half else 32} engine as {f}"
705
- )
706
- if builder.platform_has_fast_fp16 and self.args.half:
873
+ LOGGER.info(f"{prefix} building {'INT8' if int8 else 'FP' + ('16' if half else '32')} engine as {f}")
874
+ if int8:
875
+ config.set_flag(trt.BuilderFlag.INT8)
876
+ config.set_calibration_profile(profile)
877
+ config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
878
+
879
+ class EngineCalibrator(trt.IInt8Calibrator):
880
+ def __init__(
881
+ self,
882
+ dataset, # ultralytics.data.build.InfiniteDataLoader
883
+ batch: int,
884
+ cache: str = "",
885
+ ) -> None:
886
+ trt.IInt8Calibrator.__init__(self)
887
+ self.dataset = dataset
888
+ self.data_iter = iter(dataset)
889
+ self.algo = trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2
890
+ self.batch = batch
891
+ self.cache = Path(cache)
892
+
893
+ def get_algorithm(self) -> trt.CalibrationAlgoType:
894
+ """Get the calibration algorithm to use."""
895
+ return self.algo
896
+
897
+ def get_batch_size(self) -> int:
898
+ """Get the batch size to use for calibration."""
899
+ return self.batch or 1
900
+
901
+ def get_batch(self, names) -> list:
902
+ """Get the next batch to use for calibration, as a list of device memory pointers."""
903
+ try:
904
+ im0s = next(self.data_iter)["img"] / 255.0
905
+ im0s = im0s.to("cuda") if im0s.device.type == "cpu" else im0s
906
+ return [int(im0s.data_ptr())]
907
+ except StopIteration:
908
+ # Return [] or None, signal to TensorRT there is no calibration data remaining
909
+ return None
910
+
911
+ def read_calibration_cache(self) -> bytes:
912
+ """Use existing cache instead of calibrating again, otherwise, implicitly return None."""
913
+ if self.cache.exists() and self.cache.suffix == ".cache":
914
+ return self.cache.read_bytes()
915
+
916
+ def write_calibration_cache(self, cache) -> None:
917
+ """Write calibration cache to disk."""
918
+ _ = self.cache.write_bytes(cache)
919
+
920
+ # Load dataset w/ builder (for batching) and calibrate
921
+ config.int8_calibrator = EngineCalibrator(
922
+ dataset=self.get_int8_calibration_dataloader(prefix),
923
+ batch=2 * self.args.batch, # TensorRT INT8 calibration should use 2x batch size
924
+ cache=str(self.file.with_suffix(".cache")),
925
+ )
926
+
927
+ elif half:
707
928
  config.set_flag(trt.BuilderFlag.FP16)
708
929
 
930
+ # Free CUDA memory
709
931
  del self.model
932
+ gc.collect()
710
933
  torch.cuda.empty_cache()
711
934
 
712
935
  # Write file
713
- with builder.build_engine(network, config) as engine, open(f, "wb") as t:
936
+ build = builder.build_serialized_network if is_trt10 else builder.build_engine
937
+ with build(network, config) as engine, open(f, "wb") as t:
714
938
  # Metadata
715
939
  meta = json.dumps(self.metadata)
716
940
  t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
717
941
  t.write(meta.encode())
718
942
  # Model
719
- t.write(engine.serialize())
943
+ t.write(engine if is_trt10 else engine.serialize())
720
944
 
721
945
  return f, None
722
946
 
723
947
  @try_export
724
948
  def export_saved_model(self, prefix=colorstr("TensorFlow SavedModel:")):
725
- """YOLOv8 TensorFlow SavedModel export."""
949
+ """YOLO TensorFlow SavedModel export."""
726
950
  cuda = torch.cuda.is_available()
727
951
  try:
728
952
  import tensorflow as tf # noqa
729
953
  except ImportError:
730
954
  suffix = "-macos" if MACOS else "-aarch64" if ARM64 else "" if cuda else "-cpu"
731
- version = "" if ARM64 else "<=2.13.1"
955
+ version = ">=2.0.0"
732
956
  check_requirements(f"tensorflow{suffix}{version}")
733
957
  import tensorflow as tf # noqa
734
- if ARM64:
735
- check_requirements("cmake") # 'cmake' is needed to build onnxsim on aarch64
736
958
  check_requirements(
737
959
  (
960
+ "keras", # required by 'onnx2tf' package
961
+ "tf_keras", # required by 'onnx2tf' package
962
+ "sng4onnx>=1.0.1", # required by 'onnx2tf' package
963
+ "onnx_graphsurgeon>=0.3.26", # required by 'onnx2tf' package
738
964
  "onnx>=1.12.0",
739
- "onnx2tf>=1.15.4,<=1.17.5",
740
- "sng4onnx>=1.0.1",
741
- "onnxsim>=0.4.33",
742
- "onnx_graphsurgeon>=0.3.26",
743
- "tflite_support",
965
+ "onnx2tf>1.17.5,<=1.26.3",
966
+ "onnxslim>=0.1.31",
967
+ "tflite_support<=0.4.3" if IS_JETSON else "tflite_support", # fix ImportError 'GLIBCXX_3.4.29'
744
968
  "flatbuffers>=23.5.26,<100", # update old 'flatbuffers' included inside tensorflow package
745
969
  "onnxruntime-gpu" if cuda else "onnxruntime",
746
970
  ),
747
- cmds="--extra-index-url https://pypi.ngc.nvidia.com",
748
- ) # onnx_graphsurgeon only on NVIDIA
971
+ cmds="--extra-index-url https://pypi.ngc.nvidia.com", # onnx_graphsurgeon only on NVIDIA
972
+ )
749
973
 
750
974
  LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
751
975
  check_version(
752
976
  tf.__version__,
753
- "<=2.13.1",
977
+ ">=2.0.0",
754
978
  name="tensorflow",
755
979
  verbose=True,
756
980
  msg="https://github.com/ultralytics/ultralytics/issues/5161",
@@ -771,39 +995,29 @@ class Exporter:
771
995
  f_onnx, _ = self.export_onnx()
772
996
 
773
997
  # Export to TF
774
- tmp_file = f / "tmp_tflite_int8_calibration_images.npy" # int8 calibration images file
775
998
  np_data = None
776
999
  if self.args.int8:
777
- verbosity = "info"
1000
+ tmp_file = f / "tmp_tflite_int8_calibration_images.npy" # int8 calibration images file
778
1001
  if self.args.data:
779
- # Generate calibration data for integer quantization
780
- LOGGER.info(f"{prefix} collecting INT8 calibration images from 'data={self.args.data}'")
781
- data = check_det_dataset(self.args.data)
782
- dataset = YOLODataset(data["val"], data=data, imgsz=self.imgsz[0], augment=False)
783
- images = []
784
- for i, batch in enumerate(dataset):
785
- if i >= 100: # maximum number of calibration images
786
- break
787
- im = batch["img"].permute(1, 2, 0)[None] # list to nparray, CHW to BHWC
788
- images.append(im)
789
1002
  f.mkdir()
790
- images = torch.cat(images, 0).float()
791
- # mean = images.view(-1, 3).mean(0) # imagenet mean [123.675, 116.28, 103.53]
792
- # std = images.view(-1, 3).std(0) # imagenet std [58.395, 57.12, 57.375]
793
- np.save(str(tmp_file), images.numpy()) # BHWC
1003
+ images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
1004
+ images = torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz).permute(
1005
+ 0, 2, 3, 1
1006
+ )
1007
+ np.save(str(tmp_file), images.numpy().astype(np.float32)) # BHWC
794
1008
  np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
795
- else:
796
- verbosity = "error"
797
1009
 
798
1010
  LOGGER.info(f"{prefix} starting TFLite export with onnx2tf {onnx2tf.__version__}...")
799
- onnx2tf.convert(
1011
+ keras_model = onnx2tf.convert(
800
1012
  input_onnx_file_path=f_onnx,
801
1013
  output_folder_path=str(f),
802
1014
  not_use_onnxsim=True,
803
- verbosity=verbosity,
1015
+ verbosity="error", # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
804
1016
  output_integer_quantized_tflite=self.args.int8,
805
1017
  quant_type="per-tensor", # "per-tensor" (faster) or "per-channel" (slower but more accurate)
806
1018
  custom_input_op_name_np_data_path=np_data,
1019
+ disable_group_convolution=True, # for end-to-end model compatibility
1020
+ enable_batchmatmul_unfold=True, # for end-to-end model compatibility
807
1021
  )
808
1022
  yaml_save(f / "metadata.yaml", self.metadata) # add metadata.yaml
809
1023
 
@@ -819,11 +1033,11 @@ class Exporter:
819
1033
  for file in f.rglob("*.tflite"):
820
1034
  f.unlink() if "quant_with_int16_act.tflite" in str(f) else self._add_tflite_metadata(file)
821
1035
 
822
- return str(f), tf.saved_model.load(f, tags=None, options=None) # load saved_model as Keras model
1036
+ return str(f), keras_model # or keras_model = tf.saved_model.load(f, tags=None, options=None)
823
1037
 
824
1038
  @try_export
825
1039
  def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
826
- """YOLOv8 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow."""
1040
+ """YOLO TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow."""
827
1041
  import tensorflow as tf # noqa
828
1042
  from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa
829
1043
 
@@ -839,7 +1053,8 @@ class Exporter:
839
1053
 
840
1054
  @try_export
841
1055
  def export_tflite(self, keras_model, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")):
842
- """YOLOv8 TensorFlow Lite export."""
1056
+ """YOLO TensorFlow Lite export."""
1057
+ # BUG https://github.com/ultralytics/ultralytics/issues/13436
843
1058
  import tensorflow as tf # noqa
844
1059
 
845
1060
  LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
@@ -854,7 +1069,7 @@ class Exporter:
854
1069
 
855
1070
  @try_export
856
1071
  def export_edgetpu(self, tflite_model="", prefix=colorstr("Edge TPU:")):
857
- """YOLOv8 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/."""
1072
+ """YOLO Edge TPU export https://coral.ai/docs/edgetpu/models-intro/."""
858
1073
  LOGGER.warning(f"{prefix} WARNING ⚠️ Edge TPU known bug https://github.com/ultralytics/ultralytics/issues/1185")
859
1074
 
860
1075
  cmd = "edgetpu_compiler --version"
@@ -862,7 +1077,6 @@ class Exporter:
862
1077
  assert LINUX, f"export only supported on Linux. See {help_url}"
863
1078
  if subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True).returncode != 0:
864
1079
  LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
865
- sudo = subprocess.run("sudo --version >/dev/null", shell=True).returncode == 0 # sudo installed on system
866
1080
  for c in (
867
1081
  "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -",
868
1082
  'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | '
@@ -870,13 +1084,21 @@ class Exporter:
870
1084
  "sudo apt-get update",
871
1085
  "sudo apt-get install edgetpu-compiler",
872
1086
  ):
873
- subprocess.run(c if sudo else c.replace("sudo ", ""), shell=True, check=True)
1087
+ subprocess.run(c if is_sudo_available() else c.replace("sudo ", ""), shell=True, check=True)
874
1088
  ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
875
1089
 
876
1090
  LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
877
1091
  f = str(tflite_model).replace(".tflite", "_edgetpu.tflite") # Edge TPU model
878
1092
 
879
- cmd = f'edgetpu_compiler -s -d -k 10 --out_dir "{Path(f).parent}" "{tflite_model}"'
1093
+ cmd = (
1094
+ "edgetpu_compiler "
1095
+ f'--out_dir "{Path(f).parent}" '
1096
+ "--show_operations "
1097
+ "--search_delegate "
1098
+ "--delegate_search_step 30 "
1099
+ "--timeout_sec 180 "
1100
+ f'"{tflite_model}"'
1101
+ )
880
1102
  LOGGER.info(f"{prefix} running '{cmd}'")
881
1103
  subprocess.run(cmd, shell=True)
882
1104
  self._add_tflite_metadata(f)
@@ -884,7 +1106,7 @@ class Exporter:
884
1106
 
885
1107
  @try_export
886
1108
  def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
887
- """YOLOv8 TensorFlow.js export."""
1109
+ """YOLO TensorFlow.js export."""
888
1110
  check_requirements("tensorflowjs")
889
1111
  if ARM64:
890
1112
  # Fix error: `np.object` was a deprecated alias for the builtin `object` when exporting to TF.js on ARM64
@@ -914,31 +1136,160 @@ class Exporter:
914
1136
  if " " in f:
915
1137
  LOGGER.warning(f"{prefix} WARNING ⚠️ your model may not work correctly with spaces in path '{f}'.")
916
1138
 
917
- # f_json = Path(f) / 'model.json' # *.json path
918
- # with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order
919
- # subst = re.sub(
920
- # r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
921
- # r'"Identity.?.?": {"name": "Identity.?.?"}, '
922
- # r'"Identity.?.?": {"name": "Identity.?.?"}, '
923
- # r'"Identity.?.?": {"name": "Identity.?.?"}}}',
924
- # r'{"outputs": {"Identity": {"name": "Identity"}, '
925
- # r'"Identity_1": {"name": "Identity_1"}, '
926
- # r'"Identity_2": {"name": "Identity_2"}, '
927
- # r'"Identity_3": {"name": "Identity_3"}}}',
928
- # f_json.read_text(),
929
- # )
930
- # j.write(subst)
1139
+ # Add metadata
931
1140
  yaml_save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
932
1141
  return f, None
933
1142
 
1143
+ @try_export
1144
+ def export_imx(self, prefix=colorstr("IMX:")):
1145
+ """YOLO IMX export."""
1146
+ gptq = False
1147
+ assert LINUX, (
1148
+ "export only supported on Linux. See https://developer.aitrios.sony-semicon.com/en/raspberrypi-ai-camera/documentation/imx500-converter"
1149
+ )
1150
+ if getattr(self.model, "end2end", False):
1151
+ raise ValueError("IMX export is not supported for end2end models.")
1152
+ if "C2f" not in self.model.__str__():
1153
+ raise ValueError("IMX export is only supported for YOLOv8n detection models")
1154
+ check_requirements(("model-compression-toolkit==2.1.1", "sony-custom-layers==0.2.0", "tensorflow==2.12.0"))
1155
+ check_requirements("imx500-converter[pt]==3.14.3") # Separate requirements for imx500-converter
1156
+
1157
+ import model_compression_toolkit as mct
1158
+ import onnx
1159
+ from sony_custom_layers.pytorch.object_detection.nms import multiclass_nms
1160
+
1161
+ try:
1162
+ out = subprocess.run(
1163
+ ["java", "--version"], check=True, capture_output=True
1164
+ ) # Java 17 is required for imx500-converter
1165
+ if "openjdk 17" not in str(out.stdout):
1166
+ raise FileNotFoundError
1167
+ except FileNotFoundError:
1168
+ c = ["apt", "install", "-y", "openjdk-17-jdk", "openjdk-17-jre"]
1169
+ if is_sudo_available():
1170
+ c.insert(0, "sudo")
1171
+ subprocess.run(c, check=True)
1172
+
1173
+ def representative_dataset_gen(dataloader=self.get_int8_calibration_dataloader(prefix)):
1174
+ for batch in dataloader:
1175
+ img = batch["img"]
1176
+ img = img / 255.0
1177
+ yield [img]
1178
+
1179
+ tpc = mct.get_target_platform_capabilities(
1180
+ fw_name="pytorch", target_platform_name="imx500", target_platform_version="v1"
1181
+ )
1182
+
1183
+ config = mct.core.CoreConfig(
1184
+ mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=10),
1185
+ quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True),
1186
+ )
1187
+
1188
+ resource_utilization = mct.core.ResourceUtilization(weights_memory=3146176 * 0.76)
1189
+
1190
+ quant_model = (
1191
+ mct.gptq.pytorch_gradient_post_training_quantization( # Perform Gradient-Based Post Training Quantization
1192
+ model=self.model,
1193
+ representative_data_gen=representative_dataset_gen,
1194
+ target_resource_utilization=resource_utilization,
1195
+ gptq_config=mct.gptq.get_pytorch_gptq_config(n_epochs=1000, use_hessian_based_weights=False),
1196
+ core_config=config,
1197
+ target_platform_capabilities=tpc,
1198
+ )[0]
1199
+ if gptq
1200
+ else mct.ptq.pytorch_post_training_quantization( # Perform post training quantization
1201
+ in_module=self.model,
1202
+ representative_data_gen=representative_dataset_gen,
1203
+ target_resource_utilization=resource_utilization,
1204
+ core_config=config,
1205
+ target_platform_capabilities=tpc,
1206
+ )[0]
1207
+ )
1208
+
1209
+ class NMSWrapper(torch.nn.Module):
1210
+ def __init__(
1211
+ self,
1212
+ model: torch.nn.Module,
1213
+ score_threshold: float = 0.001,
1214
+ iou_threshold: float = 0.7,
1215
+ max_detections: int = 300,
1216
+ ):
1217
+ """
1218
+ Wrapping PyTorch Module with multiclass_nms layer from sony_custom_layers.
1219
+
1220
+ Args:
1221
+ model (nn.Module): Model instance.
1222
+ score_threshold (float): Score threshold for non-maximum suppression.
1223
+ iou_threshold (float): Intersection over union threshold for non-maximum suppression.
1224
+ max_detections (float): The number of detections to return.
1225
+ """
1226
+ super().__init__()
1227
+ self.model = model
1228
+ self.score_threshold = score_threshold
1229
+ self.iou_threshold = iou_threshold
1230
+ self.max_detections = max_detections
1231
+
1232
+ def forward(self, images):
1233
+ # model inference
1234
+ outputs = self.model(images)
1235
+
1236
+ boxes = outputs[0]
1237
+ scores = outputs[1]
1238
+ nms = multiclass_nms(
1239
+ boxes=boxes,
1240
+ scores=scores,
1241
+ score_threshold=self.score_threshold,
1242
+ iou_threshold=self.iou_threshold,
1243
+ max_detections=self.max_detections,
1244
+ )
1245
+ return nms
1246
+
1247
+ quant_model = NMSWrapper(
1248
+ model=quant_model,
1249
+ score_threshold=self.args.conf or 0.001,
1250
+ iou_threshold=self.args.iou,
1251
+ max_detections=self.args.max_det,
1252
+ ).to(self.device)
1253
+
1254
+ f = Path(str(self.file).replace(self.file.suffix, "_imx_model"))
1255
+ f.mkdir(exist_ok=True)
1256
+ onnx_model = f / Path(str(self.file).replace(self.file.suffix, "_imx.onnx")) # js dir
1257
+ mct.exporter.pytorch_export_model(
1258
+ model=quant_model, save_model_path=onnx_model, repr_dataset=representative_dataset_gen
1259
+ )
1260
+
1261
+ model_onnx = onnx.load(onnx_model) # load onnx model
1262
+ for k, v in self.metadata.items():
1263
+ meta = model_onnx.metadata_props.add()
1264
+ meta.key, meta.value = k, str(v)
1265
+
1266
+ onnx.save(model_onnx, onnx_model)
1267
+
1268
+ subprocess.run(
1269
+ ["imxconv-pt", "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
1270
+ check=True,
1271
+ )
1272
+
1273
+ # Needed for imx models.
1274
+ with open(f / "labels.txt", "w") as file:
1275
+ file.writelines([f"{name}\n" for _, name in self.model.names.items()])
1276
+
1277
+ return f, None
1278
+
934
1279
  def _add_tflite_metadata(self, file):
935
1280
  """Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata."""
936
- from tflite_support import flatbuffers # noqa
937
- from tflite_support import metadata as _metadata # noqa
938
- from tflite_support import metadata_schema_py_generated as _metadata_fb # noqa
1281
+ import flatbuffers
1282
+
1283
+ try:
1284
+ # TFLite Support bug https://github.com/tensorflow/tflite-support/issues/954#issuecomment-2108570845
1285
+ from tensorflow_lite_support.metadata import metadata_schema_py_generated as schema # noqa
1286
+ from tensorflow_lite_support.metadata.python import metadata # noqa
1287
+ except ImportError: # ARM64 systems may not have the 'tensorflow_lite_support' package available
1288
+ from tflite_support import metadata # noqa
1289
+ from tflite_support import metadata_schema_py_generated as schema # noqa
939
1290
 
940
1291
  # Create model info
941
- model_meta = _metadata_fb.ModelMetadataT()
1292
+ model_meta = schema.ModelMetadataT()
942
1293
  model_meta.name = self.metadata["description"]
943
1294
  model_meta.version = self.metadata["version"]
944
1295
  model_meta.author = self.metadata["author"]
@@ -949,48 +1300,48 @@ class Exporter:
949
1300
  with open(tmp_file, "w") as f:
950
1301
  f.write(str(self.metadata))
951
1302
 
952
- label_file = _metadata_fb.AssociatedFileT()
1303
+ label_file = schema.AssociatedFileT()
953
1304
  label_file.name = tmp_file.name
954
- label_file.type = _metadata_fb.AssociatedFileType.TENSOR_AXIS_LABELS
1305
+ label_file.type = schema.AssociatedFileType.TENSOR_AXIS_LABELS
955
1306
 
956
1307
  # Create input info
957
- input_meta = _metadata_fb.TensorMetadataT()
1308
+ input_meta = schema.TensorMetadataT()
958
1309
  input_meta.name = "image"
959
1310
  input_meta.description = "Input image to be detected."
960
- input_meta.content = _metadata_fb.ContentT()
961
- input_meta.content.contentProperties = _metadata_fb.ImagePropertiesT()
962
- input_meta.content.contentProperties.colorSpace = _metadata_fb.ColorSpaceType.RGB
963
- input_meta.content.contentPropertiesType = _metadata_fb.ContentProperties.ImageProperties
1311
+ input_meta.content = schema.ContentT()
1312
+ input_meta.content.contentProperties = schema.ImagePropertiesT()
1313
+ input_meta.content.contentProperties.colorSpace = schema.ColorSpaceType.RGB
1314
+ input_meta.content.contentPropertiesType = schema.ContentProperties.ImageProperties
964
1315
 
965
1316
  # Create output info
966
- output1 = _metadata_fb.TensorMetadataT()
1317
+ output1 = schema.TensorMetadataT()
967
1318
  output1.name = "output"
968
1319
  output1.description = "Coordinates of detected objects, class labels, and confidence score"
969
1320
  output1.associatedFiles = [label_file]
970
1321
  if self.model.task == "segment":
971
- output2 = _metadata_fb.TensorMetadataT()
1322
+ output2 = schema.TensorMetadataT()
972
1323
  output2.name = "output"
973
1324
  output2.description = "Mask protos"
974
1325
  output2.associatedFiles = [label_file]
975
1326
 
976
1327
  # Create subgraph info
977
- subgraph = _metadata_fb.SubGraphMetadataT()
1328
+ subgraph = schema.SubGraphMetadataT()
978
1329
  subgraph.inputTensorMetadata = [input_meta]
979
1330
  subgraph.outputTensorMetadata = [output1, output2] if self.model.task == "segment" else [output1]
980
1331
  model_meta.subgraphMetadata = [subgraph]
981
1332
 
982
1333
  b = flatbuffers.Builder(0)
983
- b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
1334
+ b.Finish(model_meta.Pack(b), metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
984
1335
  metadata_buf = b.Output()
985
1336
 
986
- populator = _metadata.MetadataPopulator.with_model_file(str(file))
1337
+ populator = metadata.MetadataPopulator.with_model_file(str(file))
987
1338
  populator.load_metadata_buffer(metadata_buf)
988
1339
  populator.load_associated_files([str(tmp_file)])
989
1340
  populator.populate()
990
1341
  tmp_file.unlink()
991
1342
 
992
1343
  def _pipeline_coreml(self, model, weights_dir=None, prefix=colorstr("CoreML Pipeline:")):
993
- """YOLOv8 CoreML pipeline."""
1344
+ """YOLO CoreML pipeline."""
994
1345
  import coremltools as ct # noqa
995
1346
 
996
1347
  LOGGER.info(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
@@ -1014,27 +1365,11 @@ class Exporter:
1014
1365
  names = self.metadata["names"]
1015
1366
  nx, ny = spec.description.input[0].type.imageType.width, spec.description.input[0].type.imageType.height
1016
1367
  _, nc = out0_shape # number of anchors, number of classes
1017
- # _, nc = out0.type.multiArrayType.shape
1018
1368
  assert len(names) == nc, f"{len(names)} names found for nc={nc}" # check
1019
1369
 
1020
1370
  # Define output shapes (missing)
1021
1371
  out0.type.multiArrayType.shape[:] = out0_shape # (3780, 80)
1022
1372
  out1.type.multiArrayType.shape[:] = out1_shape # (3780, 4)
1023
- # spec.neuralNetwork.preprocessing[0].featureName = '0'
1024
-
1025
- # Flexible input shapes
1026
- # from coremltools.models.neural_network import flexible_shape_utils
1027
- # s = [] # shapes
1028
- # s.append(flexible_shape_utils.NeuralNetworkImageSize(320, 192))
1029
- # s.append(flexible_shape_utils.NeuralNetworkImageSize(640, 384)) # (height, width)
1030
- # flexible_shape_utils.add_enumerated_image_sizes(spec, feature_name='image', sizes=s)
1031
- # r = flexible_shape_utils.NeuralNetworkImageSizeRange() # shape ranges
1032
- # r.add_height_range((192, 640))
1033
- # r.add_width_range((192, 640))
1034
- # flexible_shape_utils.update_image_size_range(spec, feature_name='image', size_range=r)
1035
-
1036
- # Print
1037
- # print(spec.description)
1038
1373
 
1039
1374
  # Model from spec
1040
1375
  model = ct.models.MLModel(spec, weights_dir=weights_dir)