dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.4.7__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.
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/METADATA +64 -74
- dgenerate_ultralytics_headless-8.4.7.dist-info/RECORD +311 -0
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/WHEEL +1 -1
- tests/__init__.py +7 -9
- tests/conftest.py +8 -15
- tests/test_cli.py +1 -1
- tests/test_cuda.py +13 -10
- tests/test_engine.py +9 -9
- tests/test_exports.py +65 -13
- tests/test_integrations.py +13 -13
- tests/test_python.py +125 -69
- tests/test_solutions.py +161 -152
- ultralytics/__init__.py +1 -1
- ultralytics/cfg/__init__.py +86 -92
- ultralytics/cfg/datasets/Argoverse.yaml +7 -6
- ultralytics/cfg/datasets/DOTAv1.5.yaml +1 -1
- ultralytics/cfg/datasets/DOTAv1.yaml +1 -1
- ultralytics/cfg/datasets/ImageNet.yaml +1 -1
- ultralytics/cfg/datasets/TT100K.yaml +346 -0
- ultralytics/cfg/datasets/VOC.yaml +15 -16
- ultralytics/cfg/datasets/african-wildlife.yaml +1 -1
- ultralytics/cfg/datasets/coco-pose.yaml +21 -0
- ultralytics/cfg/datasets/coco12-formats.yaml +101 -0
- ultralytics/cfg/datasets/coco128-seg.yaml +1 -1
- ultralytics/cfg/datasets/coco8-pose.yaml +21 -0
- ultralytics/cfg/datasets/dog-pose.yaml +28 -0
- ultralytics/cfg/datasets/dota8-multispectral.yaml +1 -1
- ultralytics/cfg/datasets/dota8.yaml +2 -2
- ultralytics/cfg/datasets/hand-keypoints.yaml +26 -2
- ultralytics/cfg/datasets/kitti.yaml +27 -0
- ultralytics/cfg/datasets/lvis.yaml +5 -5
- ultralytics/cfg/datasets/open-images-v7.yaml +1 -1
- ultralytics/cfg/datasets/tiger-pose.yaml +16 -0
- ultralytics/cfg/datasets/xView.yaml +16 -16
- ultralytics/cfg/default.yaml +4 -2
- ultralytics/cfg/models/11/yolo11-pose.yaml +1 -1
- ultralytics/cfg/models/11/yoloe-11-seg.yaml +2 -2
- ultralytics/cfg/models/11/yoloe-11.yaml +2 -2
- ultralytics/cfg/models/26/yolo26-cls.yaml +33 -0
- ultralytics/cfg/models/26/yolo26-obb.yaml +52 -0
- ultralytics/cfg/models/26/yolo26-p2.yaml +60 -0
- ultralytics/cfg/models/26/yolo26-p6.yaml +62 -0
- ultralytics/cfg/models/26/yolo26-pose.yaml +53 -0
- ultralytics/cfg/models/26/yolo26-seg.yaml +52 -0
- ultralytics/cfg/models/26/yolo26.yaml +52 -0
- ultralytics/cfg/models/26/yoloe-26-seg.yaml +53 -0
- ultralytics/cfg/models/26/yoloe-26.yaml +53 -0
- ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +1 -1
- ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +1 -1
- ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +1 -1
- ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +1 -1
- ultralytics/cfg/models/v10/yolov10b.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10l.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10m.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10n.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10s.yaml +2 -2
- ultralytics/cfg/models/v10/yolov10x.yaml +2 -2
- ultralytics/cfg/models/v3/yolov3-tiny.yaml +1 -1
- ultralytics/cfg/models/v6/yolov6.yaml +1 -1
- ultralytics/cfg/models/v8/yoloe-v8-seg.yaml +9 -6
- ultralytics/cfg/models/v8/yoloe-v8.yaml +9 -6
- ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +2 -2
- ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +2 -2
- ultralytics/cfg/models/v8/yolov8-ghost.yaml +2 -2
- ultralytics/cfg/models/v8/yolov8-obb.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-p2.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-world.yaml +1 -1
- ultralytics/cfg/models/v8/yolov8-worldv2.yaml +6 -6
- ultralytics/cfg/models/v9/yolov9s.yaml +1 -1
- ultralytics/data/__init__.py +4 -4
- ultralytics/data/annotator.py +5 -6
- ultralytics/data/augment.py +300 -475
- ultralytics/data/base.py +18 -26
- ultralytics/data/build.py +147 -25
- ultralytics/data/converter.py +108 -87
- ultralytics/data/dataset.py +47 -75
- ultralytics/data/loaders.py +42 -49
- ultralytics/data/split.py +5 -6
- ultralytics/data/split_dota.py +8 -15
- ultralytics/data/utils.py +36 -45
- ultralytics/engine/exporter.py +351 -263
- ultralytics/engine/model.py +186 -225
- ultralytics/engine/predictor.py +45 -54
- ultralytics/engine/results.py +198 -325
- ultralytics/engine/trainer.py +165 -106
- ultralytics/engine/tuner.py +41 -43
- ultralytics/engine/validator.py +55 -38
- ultralytics/hub/__init__.py +16 -19
- ultralytics/hub/auth.py +6 -12
- ultralytics/hub/google/__init__.py +7 -10
- ultralytics/hub/session.py +15 -25
- ultralytics/hub/utils.py +5 -8
- ultralytics/models/__init__.py +1 -1
- ultralytics/models/fastsam/__init__.py +1 -1
- ultralytics/models/fastsam/model.py +8 -10
- ultralytics/models/fastsam/predict.py +18 -30
- ultralytics/models/fastsam/utils.py +1 -2
- ultralytics/models/fastsam/val.py +5 -7
- ultralytics/models/nas/__init__.py +1 -1
- ultralytics/models/nas/model.py +5 -8
- ultralytics/models/nas/predict.py +7 -9
- ultralytics/models/nas/val.py +1 -2
- ultralytics/models/rtdetr/__init__.py +1 -1
- ultralytics/models/rtdetr/model.py +5 -8
- ultralytics/models/rtdetr/predict.py +15 -19
- ultralytics/models/rtdetr/train.py +10 -13
- ultralytics/models/rtdetr/val.py +21 -23
- ultralytics/models/sam/__init__.py +15 -2
- ultralytics/models/sam/amg.py +14 -20
- ultralytics/models/sam/build.py +26 -19
- ultralytics/models/sam/build_sam3.py +377 -0
- ultralytics/models/sam/model.py +29 -32
- ultralytics/models/sam/modules/blocks.py +83 -144
- ultralytics/models/sam/modules/decoders.py +19 -37
- ultralytics/models/sam/modules/encoders.py +44 -101
- ultralytics/models/sam/modules/memory_attention.py +16 -30
- ultralytics/models/sam/modules/sam.py +200 -73
- ultralytics/models/sam/modules/tiny_encoder.py +64 -83
- ultralytics/models/sam/modules/transformer.py +18 -28
- ultralytics/models/sam/modules/utils.py +174 -50
- ultralytics/models/sam/predict.py +2248 -350
- ultralytics/models/sam/sam3/__init__.py +3 -0
- ultralytics/models/sam/sam3/decoder.py +546 -0
- ultralytics/models/sam/sam3/encoder.py +529 -0
- ultralytics/models/sam/sam3/geometry_encoders.py +415 -0
- ultralytics/models/sam/sam3/maskformer_segmentation.py +286 -0
- ultralytics/models/sam/sam3/model_misc.py +199 -0
- ultralytics/models/sam/sam3/necks.py +129 -0
- ultralytics/models/sam/sam3/sam3_image.py +339 -0
- ultralytics/models/sam/sam3/text_encoder_ve.py +307 -0
- ultralytics/models/sam/sam3/vitdet.py +547 -0
- ultralytics/models/sam/sam3/vl_combiner.py +160 -0
- ultralytics/models/utils/loss.py +14 -26
- ultralytics/models/utils/ops.py +13 -17
- ultralytics/models/yolo/__init__.py +1 -1
- ultralytics/models/yolo/classify/predict.py +10 -13
- ultralytics/models/yolo/classify/train.py +12 -33
- ultralytics/models/yolo/classify/val.py +30 -29
- ultralytics/models/yolo/detect/predict.py +9 -12
- ultralytics/models/yolo/detect/train.py +17 -23
- ultralytics/models/yolo/detect/val.py +77 -59
- ultralytics/models/yolo/model.py +43 -60
- ultralytics/models/yolo/obb/predict.py +7 -16
- ultralytics/models/yolo/obb/train.py +14 -17
- ultralytics/models/yolo/obb/val.py +40 -37
- ultralytics/models/yolo/pose/__init__.py +1 -1
- ultralytics/models/yolo/pose/predict.py +7 -22
- ultralytics/models/yolo/pose/train.py +13 -16
- ultralytics/models/yolo/pose/val.py +39 -58
- ultralytics/models/yolo/segment/predict.py +17 -21
- ultralytics/models/yolo/segment/train.py +7 -10
- ultralytics/models/yolo/segment/val.py +95 -47
- ultralytics/models/yolo/world/train.py +8 -14
- ultralytics/models/yolo/world/train_world.py +11 -34
- ultralytics/models/yolo/yoloe/__init__.py +7 -7
- ultralytics/models/yolo/yoloe/predict.py +16 -23
- ultralytics/models/yolo/yoloe/train.py +36 -44
- ultralytics/models/yolo/yoloe/train_seg.py +11 -11
- ultralytics/models/yolo/yoloe/val.py +15 -20
- ultralytics/nn/__init__.py +7 -7
- ultralytics/nn/autobackend.py +159 -85
- ultralytics/nn/modules/__init__.py +68 -60
- ultralytics/nn/modules/activation.py +4 -6
- ultralytics/nn/modules/block.py +260 -224
- ultralytics/nn/modules/conv.py +52 -97
- ultralytics/nn/modules/head.py +831 -299
- ultralytics/nn/modules/transformer.py +76 -88
- ultralytics/nn/modules/utils.py +16 -21
- ultralytics/nn/tasks.py +180 -195
- ultralytics/nn/text_model.py +45 -69
- ultralytics/optim/__init__.py +5 -0
- ultralytics/optim/muon.py +338 -0
- ultralytics/solutions/__init__.py +12 -12
- ultralytics/solutions/ai_gym.py +13 -19
- ultralytics/solutions/analytics.py +15 -16
- ultralytics/solutions/config.py +6 -7
- ultralytics/solutions/distance_calculation.py +10 -13
- ultralytics/solutions/heatmap.py +8 -14
- ultralytics/solutions/instance_segmentation.py +6 -9
- ultralytics/solutions/object_blurrer.py +7 -10
- ultralytics/solutions/object_counter.py +12 -19
- ultralytics/solutions/object_cropper.py +8 -14
- ultralytics/solutions/parking_management.py +34 -32
- ultralytics/solutions/queue_management.py +10 -12
- ultralytics/solutions/region_counter.py +9 -12
- ultralytics/solutions/security_alarm.py +15 -20
- ultralytics/solutions/similarity_search.py +10 -15
- ultralytics/solutions/solutions.py +77 -76
- ultralytics/solutions/speed_estimation.py +7 -10
- ultralytics/solutions/streamlit_inference.py +2 -4
- ultralytics/solutions/templates/similarity-search.html +7 -18
- ultralytics/solutions/trackzone.py +7 -10
- ultralytics/solutions/vision_eye.py +5 -8
- ultralytics/trackers/__init__.py +1 -1
- ultralytics/trackers/basetrack.py +3 -5
- ultralytics/trackers/bot_sort.py +10 -27
- ultralytics/trackers/byte_tracker.py +21 -37
- ultralytics/trackers/track.py +4 -7
- ultralytics/trackers/utils/gmc.py +11 -22
- ultralytics/trackers/utils/kalman_filter.py +37 -48
- ultralytics/trackers/utils/matching.py +12 -15
- ultralytics/utils/__init__.py +124 -124
- ultralytics/utils/autobatch.py +2 -4
- ultralytics/utils/autodevice.py +17 -18
- ultralytics/utils/benchmarks.py +57 -71
- ultralytics/utils/callbacks/base.py +8 -10
- ultralytics/utils/callbacks/clearml.py +5 -13
- ultralytics/utils/callbacks/comet.py +32 -46
- ultralytics/utils/callbacks/dvc.py +13 -18
- ultralytics/utils/callbacks/mlflow.py +4 -5
- ultralytics/utils/callbacks/neptune.py +7 -15
- ultralytics/utils/callbacks/platform.py +423 -38
- ultralytics/utils/callbacks/raytune.py +3 -4
- ultralytics/utils/callbacks/tensorboard.py +25 -31
- ultralytics/utils/callbacks/wb.py +16 -14
- ultralytics/utils/checks.py +127 -85
- ultralytics/utils/cpu.py +3 -8
- ultralytics/utils/dist.py +9 -12
- ultralytics/utils/downloads.py +25 -33
- ultralytics/utils/errors.py +6 -14
- ultralytics/utils/events.py +2 -4
- ultralytics/utils/export/__init__.py +4 -236
- ultralytics/utils/export/engine.py +246 -0
- ultralytics/utils/export/imx.py +117 -63
- ultralytics/utils/export/tensorflow.py +231 -0
- ultralytics/utils/files.py +26 -30
- ultralytics/utils/git.py +9 -11
- ultralytics/utils/instance.py +30 -51
- ultralytics/utils/logger.py +212 -114
- ultralytics/utils/loss.py +601 -215
- ultralytics/utils/metrics.py +128 -156
- ultralytics/utils/nms.py +13 -16
- ultralytics/utils/ops.py +117 -166
- ultralytics/utils/patches.py +75 -21
- ultralytics/utils/plotting.py +75 -80
- ultralytics/utils/tal.py +125 -59
- ultralytics/utils/torch_utils.py +53 -79
- ultralytics/utils/tqdm.py +24 -21
- ultralytics/utils/triton.py +13 -19
- ultralytics/utils/tuner.py +19 -10
- dgenerate_ultralytics_headless-8.3.214.dist-info/RECORD +0 -283
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/entry_points.txt +0 -0
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/licenses/LICENSE +0 -0
- {dgenerate_ultralytics_headless-8.3.214.dist-info → dgenerate_ultralytics_headless-8.4.7.dist-info}/top_level.txt +0 -0
ultralytics/engine/exporter.py
CHANGED
|
@@ -4,36 +4,38 @@ Export a YOLO PyTorch model to other formats. TensorFlow exports authored by htt
|
|
|
4
4
|
|
|
5
5
|
Format | `format=argument` | Model
|
|
6
6
|
--- | --- | ---
|
|
7
|
-
PyTorch | - |
|
|
8
|
-
TorchScript | `torchscript` |
|
|
9
|
-
ONNX | `onnx` |
|
|
10
|
-
OpenVINO | `openvino` |
|
|
11
|
-
TensorRT | `engine` |
|
|
12
|
-
CoreML | `coreml` |
|
|
13
|
-
TensorFlow SavedModel | `saved_model` |
|
|
14
|
-
TensorFlow GraphDef | `pb` |
|
|
15
|
-
TensorFlow Lite | `tflite` |
|
|
16
|
-
TensorFlow Edge TPU | `edgetpu` |
|
|
17
|
-
TensorFlow.js | `tfjs` |
|
|
18
|
-
PaddlePaddle | `paddle` |
|
|
19
|
-
MNN | `mnn` |
|
|
20
|
-
NCNN | `ncnn` |
|
|
21
|
-
IMX | `imx` |
|
|
22
|
-
RKNN | `rknn` |
|
|
7
|
+
PyTorch | - | yolo26n.pt
|
|
8
|
+
TorchScript | `torchscript` | yolo26n.torchscript
|
|
9
|
+
ONNX | `onnx` | yolo26n.onnx
|
|
10
|
+
OpenVINO | `openvino` | yolo26n_openvino_model/
|
|
11
|
+
TensorRT | `engine` | yolo26n.engine
|
|
12
|
+
CoreML | `coreml` | yolo26n.mlpackage
|
|
13
|
+
TensorFlow SavedModel | `saved_model` | yolo26n_saved_model/
|
|
14
|
+
TensorFlow GraphDef | `pb` | yolo26n.pb
|
|
15
|
+
TensorFlow Lite | `tflite` | yolo26n.tflite
|
|
16
|
+
TensorFlow Edge TPU | `edgetpu` | yolo26n_edgetpu.tflite
|
|
17
|
+
TensorFlow.js | `tfjs` | yolo26n_web_model/
|
|
18
|
+
PaddlePaddle | `paddle` | yolo26n_paddle_model/
|
|
19
|
+
MNN | `mnn` | yolo26n.mnn
|
|
20
|
+
NCNN | `ncnn` | yolo26n_ncnn_model/
|
|
21
|
+
IMX | `imx` | yolo26n_imx_model/
|
|
22
|
+
RKNN | `rknn` | yolo26n_rknn_model/
|
|
23
|
+
ExecuTorch | `executorch` | yolo26n_executorch_model/
|
|
24
|
+
Axelera | `axelera` | yolo26n_axelera_model/
|
|
23
25
|
|
|
24
26
|
Requirements:
|
|
25
27
|
$ pip install "ultralytics[export]"
|
|
26
28
|
|
|
27
29
|
Python:
|
|
28
30
|
from ultralytics import YOLO
|
|
29
|
-
model = YOLO('
|
|
31
|
+
model = YOLO('yolo26n.pt')
|
|
30
32
|
results = model.export(format='onnx')
|
|
31
33
|
|
|
32
34
|
CLI:
|
|
33
|
-
$ yolo mode=export model=
|
|
35
|
+
$ yolo mode=export model=yolo26n.pt format=onnx
|
|
34
36
|
|
|
35
37
|
Inference:
|
|
36
|
-
$ yolo predict model=
|
|
38
|
+
$ yolo predict model=yolo26n.pt # PyTorch
|
|
37
39
|
yolo11n.torchscript # TorchScript
|
|
38
40
|
yolo11n.onnx # ONNX Runtime or OpenCV DNN with dnn=True
|
|
39
41
|
yolo11n_openvino_model # OpenVINO
|
|
@@ -48,6 +50,8 @@ Inference:
|
|
|
48
50
|
yolo11n_ncnn_model # NCNN
|
|
49
51
|
yolo11n_imx_model # IMX
|
|
50
52
|
yolo11n_rknn_model # RKNN
|
|
53
|
+
yolo11n_executorch_model # ExecuTorch
|
|
54
|
+
yolo11n_axelera_model # Axelera
|
|
51
55
|
|
|
52
56
|
TensorFlow.js:
|
|
53
57
|
$ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
|
|
@@ -62,7 +66,6 @@ import re
|
|
|
62
66
|
import shutil
|
|
63
67
|
import subprocess
|
|
64
68
|
import time
|
|
65
|
-
import warnings
|
|
66
69
|
from copy import deepcopy
|
|
67
70
|
from datetime import datetime
|
|
68
71
|
from pathlib import Path
|
|
@@ -82,13 +85,17 @@ from ultralytics.utils import (
|
|
|
82
85
|
ARM64,
|
|
83
86
|
DEFAULT_CFG,
|
|
84
87
|
IS_COLAB,
|
|
88
|
+
IS_DEBIAN_BOOKWORM,
|
|
89
|
+
IS_DEBIAN_TRIXIE,
|
|
90
|
+
IS_DOCKER,
|
|
85
91
|
IS_JETSON,
|
|
92
|
+
IS_RASPBERRYPI,
|
|
93
|
+
IS_UBUNTU,
|
|
86
94
|
LINUX,
|
|
87
95
|
LOGGER,
|
|
88
96
|
MACOS,
|
|
89
97
|
MACOS_VERSION,
|
|
90
98
|
RKNN_CHIPS,
|
|
91
|
-
ROOT,
|
|
92
99
|
SETTINGS,
|
|
93
100
|
TORCH_VERSION,
|
|
94
101
|
WINDOWS,
|
|
@@ -98,21 +105,38 @@ from ultralytics.utils import (
|
|
|
98
105
|
get_default_args,
|
|
99
106
|
)
|
|
100
107
|
from ultralytics.utils.checks import (
|
|
108
|
+
IS_PYTHON_3_10,
|
|
109
|
+
IS_PYTHON_MINIMUM_3_9,
|
|
110
|
+
check_apt_requirements,
|
|
101
111
|
check_imgsz,
|
|
102
|
-
check_is_path_safe,
|
|
103
112
|
check_requirements,
|
|
104
113
|
check_version,
|
|
105
114
|
is_intel,
|
|
106
115
|
is_sudo_available,
|
|
107
116
|
)
|
|
108
|
-
from ultralytics.utils.
|
|
109
|
-
|
|
110
|
-
|
|
117
|
+
from ultralytics.utils.export import (
|
|
118
|
+
keras2pb,
|
|
119
|
+
onnx2engine,
|
|
120
|
+
onnx2saved_model,
|
|
121
|
+
pb2tfjs,
|
|
122
|
+
tflite2edgetpu,
|
|
123
|
+
torch2imx,
|
|
124
|
+
torch2onnx,
|
|
125
|
+
)
|
|
126
|
+
from ultralytics.utils.files import file_size
|
|
111
127
|
from ultralytics.utils.metrics import batch_probiou
|
|
112
128
|
from ultralytics.utils.nms import TorchNMS
|
|
113
129
|
from ultralytics.utils.ops import Profile
|
|
114
130
|
from ultralytics.utils.patches import arange_patch
|
|
115
|
-
from ultralytics.utils.torch_utils import
|
|
131
|
+
from ultralytics.utils.torch_utils import (
|
|
132
|
+
TORCH_1_10,
|
|
133
|
+
TORCH_1_11,
|
|
134
|
+
TORCH_1_13,
|
|
135
|
+
TORCH_2_1,
|
|
136
|
+
TORCH_2_4,
|
|
137
|
+
TORCH_2_9,
|
|
138
|
+
select_device,
|
|
139
|
+
)
|
|
116
140
|
|
|
117
141
|
|
|
118
142
|
def export_formats():
|
|
@@ -148,18 +172,20 @@ def export_formats():
|
|
|
148
172
|
["NCNN", "ncnn", "_ncnn_model", True, True, ["batch", "half"]],
|
|
149
173
|
["IMX", "imx", "_imx_model", True, True, ["int8", "fraction", "nms"]],
|
|
150
174
|
["RKNN", "rknn", "_rknn_model", False, False, ["batch", "name"]],
|
|
175
|
+
["ExecuTorch", "executorch", "_executorch_model", True, False, ["batch"]],
|
|
176
|
+
["Axelera", "axelera", "_axelera_model", False, False, ["batch", "int8", "fraction"]],
|
|
151
177
|
]
|
|
152
178
|
return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU", "Arguments"], zip(*x)))
|
|
153
179
|
|
|
154
180
|
|
|
155
181
|
def best_onnx_opset(onnx, cuda=False) -> int:
|
|
156
182
|
"""Return max ONNX opset for this torch version with ONNX fallback."""
|
|
157
|
-
version = ".".join(TORCH_VERSION.split(".")[:2])
|
|
158
183
|
if TORCH_2_4: # _constants.ONNX_MAX_OPSET first defined in torch 1.13
|
|
159
184
|
opset = torch.onnx.utils._constants.ONNX_MAX_OPSET - 1 # use second-latest version for safety
|
|
160
185
|
if cuda:
|
|
161
186
|
opset -= 2 # fix CUDA ONNXRuntime NMS squeeze op errors
|
|
162
187
|
else:
|
|
188
|
+
version = ".".join(TORCH_VERSION.split(".")[:2])
|
|
163
189
|
opset = {
|
|
164
190
|
"1.8": 12,
|
|
165
191
|
"1.9": 12,
|
|
@@ -181,8 +207,7 @@ def best_onnx_opset(onnx, cuda=False) -> int:
|
|
|
181
207
|
|
|
182
208
|
|
|
183
209
|
def validate_args(format, passed_args, valid_args):
|
|
184
|
-
"""
|
|
185
|
-
Validate arguments based on the export format.
|
|
210
|
+
"""Validate arguments based on the export format.
|
|
186
211
|
|
|
187
212
|
Args:
|
|
188
213
|
format (str): The export format.
|
|
@@ -203,15 +228,6 @@ def validate_args(format, passed_args, valid_args):
|
|
|
203
228
|
assert arg in valid_args, f"ERROR ❌️ argument '{arg}' is not supported for format='{format}'"
|
|
204
229
|
|
|
205
230
|
|
|
206
|
-
def gd_outputs(gd):
|
|
207
|
-
"""Return TensorFlow GraphDef model output node names."""
|
|
208
|
-
name_list, input_list = [], []
|
|
209
|
-
for node in gd.node: # tensorflow.core.framework.node_def_pb2.NodeDef
|
|
210
|
-
name_list.append(node.name)
|
|
211
|
-
input_list.extend(node.input)
|
|
212
|
-
return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))
|
|
213
|
-
|
|
214
|
-
|
|
215
231
|
def try_export(inner_func):
|
|
216
232
|
"""YOLO export decorator, i.e. @try_export."""
|
|
217
233
|
inner_args = get_default_args(inner_func)
|
|
@@ -236,8 +252,7 @@ def try_export(inner_func):
|
|
|
236
252
|
|
|
237
253
|
|
|
238
254
|
class Exporter:
|
|
239
|
-
"""
|
|
240
|
-
A class for exporting YOLO models to various formats.
|
|
255
|
+
"""A class for exporting YOLO models to various formats.
|
|
241
256
|
|
|
242
257
|
This class provides functionality to export YOLO models to different formats including ONNX, TensorRT, CoreML,
|
|
243
258
|
TensorFlow, and others. It handles format validation, device selection, model preparation, and the actual export
|
|
@@ -287,8 +302,7 @@ class Exporter:
|
|
|
287
302
|
"""
|
|
288
303
|
|
|
289
304
|
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
|
|
290
|
-
"""
|
|
291
|
-
Initialize the Exporter class.
|
|
305
|
+
"""Initialize the Exporter class.
|
|
292
306
|
|
|
293
307
|
Args:
|
|
294
308
|
cfg (str, optional): Path to a configuration file.
|
|
@@ -300,7 +314,11 @@ class Exporter:
|
|
|
300
314
|
callbacks.add_integration_callbacks(self)
|
|
301
315
|
|
|
302
316
|
def __call__(self, model=None) -> str:
|
|
303
|
-
"""
|
|
317
|
+
"""Export a model and return the final exported path as a string.
|
|
318
|
+
|
|
319
|
+
Returns:
|
|
320
|
+
(str): Path to the exported file or directory (the last export artifact).
|
|
321
|
+
"""
|
|
304
322
|
t = time.time()
|
|
305
323
|
fmt = self.args.format.lower() # to lowercase
|
|
306
324
|
if fmt in {"tensorrt", "trt"}: # 'engine' aliases
|
|
@@ -322,9 +340,25 @@ class Exporter:
|
|
|
322
340
|
flags = [x == fmt for x in fmts]
|
|
323
341
|
if sum(flags) != 1:
|
|
324
342
|
raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
|
|
325
|
-
(
|
|
326
|
-
|
|
327
|
-
|
|
343
|
+
(
|
|
344
|
+
jit,
|
|
345
|
+
onnx,
|
|
346
|
+
xml,
|
|
347
|
+
engine,
|
|
348
|
+
coreml,
|
|
349
|
+
saved_model,
|
|
350
|
+
pb,
|
|
351
|
+
tflite,
|
|
352
|
+
edgetpu,
|
|
353
|
+
tfjs,
|
|
354
|
+
paddle,
|
|
355
|
+
mnn,
|
|
356
|
+
ncnn,
|
|
357
|
+
imx,
|
|
358
|
+
rknn,
|
|
359
|
+
executorch,
|
|
360
|
+
axelera,
|
|
361
|
+
) = flags # export booleans
|
|
328
362
|
|
|
329
363
|
is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
|
|
330
364
|
|
|
@@ -334,9 +368,10 @@ class Exporter:
|
|
|
334
368
|
LOGGER.warning("TensorRT requires GPU export, automatically assigning device=0")
|
|
335
369
|
self.args.device = "0"
|
|
336
370
|
if engine and "dla" in str(self.args.device): # convert int/list to str first
|
|
337
|
-
|
|
371
|
+
device_str = str(self.args.device)
|
|
372
|
+
dla = device_str.rsplit(":", 1)[-1]
|
|
338
373
|
self.args.device = "0" # update device to "0"
|
|
339
|
-
assert dla in {"0", "1"}, f"Expected
|
|
374
|
+
assert dla in {"0", "1"}, f"Expected device 'dla:0' or 'dla:1', but got {device_str}."
|
|
340
375
|
if imx and self.args.device is None and torch.cuda.is_available():
|
|
341
376
|
LOGGER.warning("Exporting on CPU while CUDA is available, setting device=0 for faster export on GPU.")
|
|
342
377
|
self.args.device = "0" # update device to "0"
|
|
@@ -345,23 +380,37 @@ class Exporter:
|
|
|
345
380
|
# Argument compatibility checks
|
|
346
381
|
fmt_keys = fmts_dict["Arguments"][flags.index(True) + 1]
|
|
347
382
|
validate_args(fmt, self.args, fmt_keys)
|
|
383
|
+
if axelera:
|
|
384
|
+
if not IS_PYTHON_3_10:
|
|
385
|
+
raise SystemError("Axelera export only supported on Python 3.10.")
|
|
386
|
+
if not self.args.int8:
|
|
387
|
+
LOGGER.warning("Setting int8=True for Axelera mixed-precision export.")
|
|
388
|
+
self.args.int8 = True
|
|
389
|
+
if model.task not in {"detect"}:
|
|
390
|
+
raise ValueError("Axelera export only supported for detection models.")
|
|
391
|
+
if not self.args.data:
|
|
392
|
+
self.args.data = "coco128.yaml" # Axelera default to coco128.yaml
|
|
348
393
|
if imx:
|
|
349
394
|
if not self.args.int8:
|
|
350
395
|
LOGGER.warning("IMX export requires int8=True, setting int8=True.")
|
|
351
396
|
self.args.int8 = True
|
|
352
|
-
if not self.args.nms:
|
|
397
|
+
if not self.args.nms and model.task in {"detect", "pose", "segment"}:
|
|
353
398
|
LOGGER.warning("IMX export requires nms=True, setting nms=True.")
|
|
354
399
|
self.args.nms = True
|
|
355
|
-
if model.task not in {"detect", "pose"}:
|
|
356
|
-
raise ValueError(
|
|
400
|
+
if model.task not in {"detect", "pose", "classify", "segment"}:
|
|
401
|
+
raise ValueError(
|
|
402
|
+
"IMX export only supported for detection, pose estimation, classification, and segmentation models."
|
|
403
|
+
)
|
|
357
404
|
if not hasattr(model, "names"):
|
|
358
405
|
model.names = default_class_names()
|
|
359
406
|
model.names = check_class_names(model.names)
|
|
360
407
|
if self.args.half and self.args.int8:
|
|
361
408
|
LOGGER.warning("half=True and int8=True are mutually exclusive, setting half=False.")
|
|
362
409
|
self.args.half = False
|
|
363
|
-
if self.args.half and
|
|
364
|
-
LOGGER.warning(
|
|
410
|
+
if self.args.half and jit and self.device.type == "cpu":
|
|
411
|
+
LOGGER.warning(
|
|
412
|
+
"half=True only compatible with GPU export for TorchScript, i.e. use device=0, setting half=False."
|
|
413
|
+
)
|
|
365
414
|
self.args.half = False
|
|
366
415
|
self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2) # check image size
|
|
367
416
|
if self.args.optimize:
|
|
@@ -378,14 +427,12 @@ class Exporter:
|
|
|
378
427
|
assert self.args.name in RKNN_CHIPS, (
|
|
379
428
|
f"Invalid processor name '{self.args.name}' for Rockchip RKNN export. Valid names are {RKNN_CHIPS}."
|
|
380
429
|
)
|
|
381
|
-
if self.args.int8 and tflite:
|
|
382
|
-
assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
|
|
383
430
|
if self.args.nms:
|
|
384
431
|
assert not isinstance(model, ClassificationModel), "'nms=True' is not valid for classification models."
|
|
385
432
|
assert not tflite or not ARM64 or not LINUX, "TFLite export with NMS unsupported on ARM64 Linux"
|
|
386
433
|
assert not is_tf_format or TORCH_1_13, "TensorFlow exports with NMS require torch>=1.13"
|
|
387
434
|
assert not onnx or TORCH_1_13, "ONNX export with NMS requires torch>=1.13"
|
|
388
|
-
if getattr(model, "end2end", False):
|
|
435
|
+
if getattr(model, "end2end", False) or isinstance(model.model[-1], RTDETRDecoder):
|
|
389
436
|
LOGGER.warning("'nms=True' is not available for end2end models. Forcing 'nms=False'.")
|
|
390
437
|
self.args.nms = False
|
|
391
438
|
self.args.conf = self.args.conf or 0.25 # set conf default value for nms export
|
|
@@ -416,6 +463,9 @@ class Exporter:
|
|
|
416
463
|
)
|
|
417
464
|
if tfjs and (ARM64 and LINUX):
|
|
418
465
|
raise SystemError("TF.js exports are not currently supported on ARM64 Linux")
|
|
466
|
+
if ncnn and hasattr(model.model[-1], "one2one_cv2"):
|
|
467
|
+
del model.model[-1].one2one_cv2 # Disable end2end branch for NCNN export as it does not support topk
|
|
468
|
+
LOGGER.warning("NCNN export does not support end2end models, disabling end2end branch.")
|
|
419
469
|
# Recommend OpenVINO if export and Intel CPU
|
|
420
470
|
if SETTINGS.get("openvino_msg"):
|
|
421
471
|
if is_intel():
|
|
@@ -445,6 +495,10 @@ class Exporter:
|
|
|
445
495
|
from ultralytics.utils.export.imx import FXModel
|
|
446
496
|
|
|
447
497
|
model = FXModel(model, self.imgsz)
|
|
498
|
+
if tflite or edgetpu:
|
|
499
|
+
from ultralytics.utils.export.tensorflow import tf_wrapper
|
|
500
|
+
|
|
501
|
+
model = tf_wrapper(model)
|
|
448
502
|
for m in model.modules():
|
|
449
503
|
if isinstance(m, Classify):
|
|
450
504
|
m.export = True
|
|
@@ -452,8 +506,11 @@ class Exporter:
|
|
|
452
506
|
m.dynamic = self.args.dynamic
|
|
453
507
|
m.export = True
|
|
454
508
|
m.format = self.args.format
|
|
455
|
-
|
|
509
|
+
# Clamp max_det to anchor count for small image sizes (required for TensorRT compatibility)
|
|
510
|
+
anchors = sum(int(self.imgsz[0] / s) * int(self.imgsz[1] / s) for s in model.stride.tolist())
|
|
511
|
+
m.max_det = min(self.args.max_det, anchors)
|
|
456
512
|
m.xyxy = self.args.nms and not coreml
|
|
513
|
+
m.shape = None # reset cached shape for new export input size
|
|
457
514
|
if hasattr(model, "pe") and hasattr(m, "fuse"): # for YOLOE models
|
|
458
515
|
m.fuse(model.pe.to(self.device))
|
|
459
516
|
elif isinstance(m, C2f) and not is_tf_format:
|
|
@@ -466,11 +523,6 @@ class Exporter:
|
|
|
466
523
|
if self.args.half and (onnx or jit) and self.device.type != "cpu":
|
|
467
524
|
im, model = im.half(), model.half() # to FP16
|
|
468
525
|
|
|
469
|
-
# Filter warnings
|
|
470
|
-
warnings.filterwarnings("ignore", category=torch.jit.TracerWarning) # suppress TracerWarning
|
|
471
|
-
warnings.filterwarnings("ignore", category=UserWarning) # suppress shape prim::Constant missing ONNX warning
|
|
472
|
-
warnings.filterwarnings("ignore", category=DeprecationWarning) # suppress CoreML np.bool deprecation warning
|
|
473
|
-
|
|
474
526
|
# Assign
|
|
475
527
|
self.im = im
|
|
476
528
|
self.model = model
|
|
@@ -502,6 +554,10 @@ class Exporter:
|
|
|
502
554
|
self.metadata["dla"] = dla # make sure `AutoBackend` uses correct dla device if it has one
|
|
503
555
|
if model.task == "pose":
|
|
504
556
|
self.metadata["kpt_shape"] = model.model[-1].kpt_shape
|
|
557
|
+
if hasattr(model, "kpt_names"):
|
|
558
|
+
self.metadata["kpt_names"] = model.kpt_names
|
|
559
|
+
if getattr(model.model[-1], "end2end", False):
|
|
560
|
+
self.metadata["end2end"] = True
|
|
505
561
|
|
|
506
562
|
LOGGER.info(
|
|
507
563
|
f"\n{colorstr('PyTorch:')} starting from '{file}' with input shape {tuple(im.shape)} BCHW and "
|
|
@@ -514,7 +570,7 @@ class Exporter:
|
|
|
514
570
|
f[0] = self.export_torchscript()
|
|
515
571
|
if engine: # TensorRT required before ONNX
|
|
516
572
|
f[1] = self.export_engine(dla=dla)
|
|
517
|
-
if onnx
|
|
573
|
+
if onnx: # ONNX
|
|
518
574
|
f[2] = self.export_onnx()
|
|
519
575
|
if xml: # OpenVINO
|
|
520
576
|
f[3] = self.export_openvino()
|
|
@@ -541,6 +597,10 @@ class Exporter:
|
|
|
541
597
|
f[13] = self.export_imx()
|
|
542
598
|
if rknn:
|
|
543
599
|
f[14] = self.export_rknn()
|
|
600
|
+
if executorch:
|
|
601
|
+
f[15] = self.export_executorch()
|
|
602
|
+
if axelera:
|
|
603
|
+
f[16] = self.export_axelera()
|
|
544
604
|
|
|
545
605
|
# Finish
|
|
546
606
|
f = [str(x) for x in f if x] # filter out '' and None
|
|
@@ -554,18 +614,17 @@ class Exporter:
|
|
|
554
614
|
f"work. Use export 'imgsz={max(self.imgsz)}' if val is required."
|
|
555
615
|
)
|
|
556
616
|
imgsz = self.imgsz[0] if square else str(self.imgsz)[1:-1].replace(" ", "")
|
|
557
|
-
predict_data = f"data={data}" if model.task == "segment" and pb else ""
|
|
558
617
|
q = "int8" if self.args.int8 else "half" if self.args.half else "" # quantization
|
|
559
618
|
LOGGER.info(
|
|
560
619
|
f"\nExport complete ({time.time() - t:.1f}s)"
|
|
561
620
|
f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
|
|
562
|
-
f"\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {q}
|
|
621
|
+
f"\nPredict: yolo predict task={model.task} model={f} imgsz={imgsz} {q}"
|
|
563
622
|
f"\nValidate: yolo val task={model.task} model={f} imgsz={imgsz} data={data} {q} {s}"
|
|
564
623
|
f"\nVisualize: https://netron.app"
|
|
565
624
|
)
|
|
566
625
|
|
|
567
626
|
self.run_callbacks("on_export_end")
|
|
568
|
-
return f #
|
|
627
|
+
return f # path to final export artifact
|
|
569
628
|
|
|
570
629
|
def get_int8_calibration_dataloader(self, prefix=""):
|
|
571
630
|
"""Build and return a dataloader for calibration of INT8 models."""
|
|
@@ -586,7 +645,9 @@ class Exporter:
|
|
|
586
645
|
f"The calibration dataset ({n} images) must have at least as many images as the batch size "
|
|
587
646
|
f"('batch={self.args.batch}')."
|
|
588
647
|
)
|
|
589
|
-
elif n <
|
|
648
|
+
elif self.args.format == "axelera" and n < 100:
|
|
649
|
+
LOGGER.warning(f"{prefix} >100 images required for Axelera calibration, found {n} images.")
|
|
650
|
+
elif self.args.format != "axelera" and n < 300:
|
|
590
651
|
LOGGER.warning(f"{prefix} >300 images recommended for INT8 calibration, found {n} images.")
|
|
591
652
|
return build_dataloader(dataset, batch=self.args.batch, workers=0, drop_last=True) # required for batch loading
|
|
592
653
|
|
|
@@ -610,11 +671,11 @@ class Exporter:
|
|
|
610
671
|
@try_export
|
|
611
672
|
def export_onnx(self, prefix=colorstr("ONNX:")):
|
|
612
673
|
"""Export YOLO model to ONNX format."""
|
|
613
|
-
requirements = ["onnx>=1.12.0"]
|
|
674
|
+
requirements = ["onnx>=1.12.0,<2.0.0"]
|
|
614
675
|
if self.args.simplify:
|
|
615
676
|
requirements += ["onnxslim>=0.1.71", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
|
|
616
677
|
check_requirements(requirements)
|
|
617
|
-
import onnx
|
|
678
|
+
import onnx
|
|
618
679
|
|
|
619
680
|
opset = self.args.opset or best_onnx_opset(onnx, cuda="cuda" in self.device.type)
|
|
620
681
|
LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__} opset {opset}...")
|
|
@@ -622,7 +683,7 @@ class Exporter:
|
|
|
622
683
|
assert TORCH_1_13, f"'nms=True' ONNX export requires torch>=1.13 (found torch=={TORCH_VERSION})"
|
|
623
684
|
|
|
624
685
|
f = str(self.file.with_suffix(".onnx"))
|
|
625
|
-
output_names = ["output0", "output1"] if
|
|
686
|
+
output_names = ["output0", "output1"] if self.model.task == "segment" else ["output0"]
|
|
626
687
|
dynamic = self.args.dynamic
|
|
627
688
|
if dynamic:
|
|
628
689
|
dynamic = {"images": {0: "batch", 2: "height", 3: "width"}} # shape(1,3,640,640)
|
|
@@ -671,6 +732,16 @@ class Exporter:
|
|
|
671
732
|
LOGGER.info(f"{prefix} limiting IR version {model_onnx.ir_version} to 10 for ONNXRuntime compatibility...")
|
|
672
733
|
model_onnx.ir_version = 10
|
|
673
734
|
|
|
735
|
+
# FP16 conversion for CPU export (GPU exports are already FP16 from model.half() during tracing)
|
|
736
|
+
if self.args.half and self.args.format == "onnx" and self.device.type == "cpu":
|
|
737
|
+
try:
|
|
738
|
+
from onnxruntime.transformers import float16
|
|
739
|
+
|
|
740
|
+
LOGGER.info(f"{prefix} converting to FP16...")
|
|
741
|
+
model_onnx = float16.convert_float_to_float16(model_onnx, keep_io_types=True)
|
|
742
|
+
except Exception as e:
|
|
743
|
+
LOGGER.warning(f"{prefix} FP16 conversion failure: {e}")
|
|
744
|
+
|
|
674
745
|
onnx.save(model_onnx, f)
|
|
675
746
|
return f
|
|
676
747
|
|
|
@@ -711,13 +782,6 @@ class Exporter:
|
|
|
711
782
|
check_requirements("nncf>=2.14.0")
|
|
712
783
|
import nncf
|
|
713
784
|
|
|
714
|
-
def transform_fn(data_item) -> np.ndarray:
|
|
715
|
-
"""Quantization transform function."""
|
|
716
|
-
data_item: torch.Tensor = data_item["img"] if isinstance(data_item, dict) else data_item
|
|
717
|
-
assert data_item.dtype == torch.uint8, "Input image must be uint8 for the quantization preprocessing"
|
|
718
|
-
im = data_item.numpy().astype(np.float32) / 255.0 # uint8 to fp16/32 and 0-255 to 0.0-1.0
|
|
719
|
-
return np.expand_dims(im, 0) if im.ndim == 3 else im
|
|
720
|
-
|
|
721
785
|
# Generate calibration data for integer quantization
|
|
722
786
|
ignored_scope = None
|
|
723
787
|
if isinstance(self.model.model[-1], Detect):
|
|
@@ -729,14 +793,13 @@ class Exporter:
|
|
|
729
793
|
f".*{head_module_name}/.*/Sub*",
|
|
730
794
|
f".*{head_module_name}/.*/Mul*",
|
|
731
795
|
f".*{head_module_name}/.*/Div*",
|
|
732
|
-
f".*{head_module_name}\\.dfl.*",
|
|
733
796
|
],
|
|
734
797
|
types=["Sigmoid"],
|
|
735
798
|
)
|
|
736
799
|
|
|
737
800
|
quantized_ov_model = nncf.quantize(
|
|
738
801
|
model=ov_model,
|
|
739
|
-
calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix),
|
|
802
|
+
calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix), self._transform_fn),
|
|
740
803
|
preset=nncf.QuantizationPreset.MIXED,
|
|
741
804
|
ignored_scope=ignored_scope,
|
|
742
805
|
)
|
|
@@ -755,16 +818,16 @@ class Exporter:
|
|
|
755
818
|
assert not IS_JETSON, "Jetson Paddle exports not supported yet"
|
|
756
819
|
check_requirements(
|
|
757
820
|
(
|
|
758
|
-
"paddlepaddle-gpu"
|
|
821
|
+
"paddlepaddle-gpu>=3.0.0,!=3.3.0" # exclude 3.3.0 https://github.com/PaddlePaddle/Paddle/issues/77340
|
|
759
822
|
if torch.cuda.is_available()
|
|
760
823
|
else "paddlepaddle==3.0.0" # pin 3.0.0 for ARM64
|
|
761
824
|
if ARM64
|
|
762
|
-
else "paddlepaddle>=3.0.0",
|
|
825
|
+
else "paddlepaddle>=3.0.0,!=3.3.0", # exclude 3.3.0 https://github.com/PaddlePaddle/Paddle/issues/77340
|
|
763
826
|
"x2paddle",
|
|
764
827
|
)
|
|
765
828
|
)
|
|
766
|
-
import x2paddle
|
|
767
|
-
from x2paddle.convert import pytorch2paddle
|
|
829
|
+
import x2paddle
|
|
830
|
+
from x2paddle.convert import pytorch2paddle
|
|
768
831
|
|
|
769
832
|
LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
|
|
770
833
|
f = str(self.file).replace(self.file.suffix, f"_paddle_model{os.sep}")
|
|
@@ -776,10 +839,11 @@ class Exporter:
|
|
|
776
839
|
@try_export
|
|
777
840
|
def export_mnn(self, prefix=colorstr("MNN:")):
|
|
778
841
|
"""Export YOLO model to MNN format using MNN https://github.com/alibaba/MNN."""
|
|
842
|
+
assert TORCH_1_10, "MNN export requires torch>=1.10.0 to avoid segmentation faults"
|
|
779
843
|
f_onnx = self.export_onnx() # get onnx model first
|
|
780
844
|
|
|
781
845
|
check_requirements("MNN>=2.9.6")
|
|
782
|
-
import MNN
|
|
846
|
+
import MNN
|
|
783
847
|
from MNN.tools import mnnconvert
|
|
784
848
|
|
|
785
849
|
# Setup and checks
|
|
@@ -802,65 +866,31 @@ class Exporter:
|
|
|
802
866
|
def export_ncnn(self, prefix=colorstr("NCNN:")):
|
|
803
867
|
"""Export YOLO model to NCNN format using PNNX https://github.com/pnnx/pnnx."""
|
|
804
868
|
check_requirements("ncnn", cmds="--no-deps") # no deps to avoid installing opencv-python
|
|
805
|
-
|
|
869
|
+
check_requirements("pnnx")
|
|
870
|
+
import ncnn
|
|
871
|
+
import pnnx
|
|
806
872
|
|
|
807
|
-
LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__}...")
|
|
873
|
+
LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__} and PNNX {pnnx.__version__}...")
|
|
808
874
|
f = Path(str(self.file).replace(self.file.suffix, f"_ncnn_model{os.sep}"))
|
|
809
|
-
f_onnx = self.file.with_suffix(".onnx")
|
|
810
875
|
|
|
811
|
-
|
|
812
|
-
|
|
813
|
-
|
|
814
|
-
|
|
815
|
-
|
|
816
|
-
|
|
817
|
-
|
|
818
|
-
)
|
|
819
|
-
|
|
820
|
-
|
|
821
|
-
|
|
822
|
-
|
|
823
|
-
|
|
824
|
-
|
|
825
|
-
except Exception as e:
|
|
826
|
-
release = "20250930"
|
|
827
|
-
asset = f"pnnx-{release}-{system}.zip"
|
|
828
|
-
LOGGER.warning(f"{prefix} PNNX GitHub assets not found: {e}, using default {asset}")
|
|
829
|
-
unzip_dir = safe_download(f"https://github.com/pnnx/pnnx/releases/download/{release}/{asset}", delete=True)
|
|
830
|
-
if check_is_path_safe(Path.cwd(), unzip_dir): # avoid path traversal security vulnerability
|
|
831
|
-
shutil.move(src=unzip_dir / name, dst=pnnx) # move binary to ROOT
|
|
832
|
-
pnnx.chmod(0o777) # set read, write, and execute permissions for everyone
|
|
833
|
-
shutil.rmtree(unzip_dir) # delete unzip dir
|
|
834
|
-
|
|
835
|
-
ncnn_args = [
|
|
836
|
-
f"ncnnparam={f / 'model.ncnn.param'}",
|
|
837
|
-
f"ncnnbin={f / 'model.ncnn.bin'}",
|
|
838
|
-
f"ncnnpy={f / 'model_ncnn.py'}",
|
|
839
|
-
]
|
|
840
|
-
|
|
841
|
-
pnnx_args = [
|
|
842
|
-
f"pnnxparam={f / 'model.pnnx.param'}",
|
|
843
|
-
f"pnnxbin={f / 'model.pnnx.bin'}",
|
|
844
|
-
f"pnnxpy={f / 'model_pnnx.py'}",
|
|
845
|
-
f"pnnxonnx={f / 'model.pnnx.onnx'}",
|
|
846
|
-
]
|
|
847
|
-
|
|
848
|
-
cmd = [
|
|
849
|
-
str(pnnx),
|
|
850
|
-
str(f_onnx),
|
|
851
|
-
*ncnn_args,
|
|
852
|
-
*pnnx_args,
|
|
853
|
-
f"fp16={int(self.args.half)}",
|
|
854
|
-
f"device={self.device.type}",
|
|
855
|
-
f'inputshape="{[self.args.batch, 3, *self.imgsz]}"',
|
|
856
|
-
]
|
|
876
|
+
ncnn_args = dict(
|
|
877
|
+
ncnnparam=(f / "model.ncnn.param").as_posix(),
|
|
878
|
+
ncnnbin=(f / "model.ncnn.bin").as_posix(),
|
|
879
|
+
ncnnpy=(f / "model_ncnn.py").as_posix(),
|
|
880
|
+
)
|
|
881
|
+
|
|
882
|
+
pnnx_args = dict(
|
|
883
|
+
ptpath=(f / "model.pt").as_posix(),
|
|
884
|
+
pnnxparam=(f / "model.pnnx.param").as_posix(),
|
|
885
|
+
pnnxbin=(f / "model.pnnx.bin").as_posix(),
|
|
886
|
+
pnnxpy=(f / "model_pnnx.py").as_posix(),
|
|
887
|
+
pnnxonnx=(f / "model.pnnx.onnx").as_posix(),
|
|
888
|
+
)
|
|
889
|
+
|
|
857
890
|
f.mkdir(exist_ok=True) # make ncnn_model directory
|
|
858
|
-
|
|
859
|
-
subprocess.run(cmd, check=True)
|
|
891
|
+
pnnx.export(self.model, inputs=self.im, **ncnn_args, **pnnx_args, fp16=self.args.half, device=self.device.type)
|
|
860
892
|
|
|
861
|
-
|
|
862
|
-
pnnx_files = [x.rsplit("=", 1)[-1] for x in pnnx_args]
|
|
863
|
-
for f_debug in ("debug.bin", "debug.param", "debug2.bin", "debug2.param", *pnnx_files):
|
|
893
|
+
for f_debug in ("debug.bin", "debug.param", "debug2.bin", "debug2.param", *pnnx_args.values()):
|
|
864
894
|
Path(f_debug).unlink(missing_ok=True)
|
|
865
895
|
|
|
866
896
|
YAML.save(f / "metadata.yaml", self.metadata) # add metadata.yaml
|
|
@@ -870,8 +900,10 @@ class Exporter:
|
|
|
870
900
|
def export_coreml(self, prefix=colorstr("CoreML:")):
|
|
871
901
|
"""Export YOLO model to CoreML format."""
|
|
872
902
|
mlmodel = self.args.format.lower() == "mlmodel" # legacy *.mlmodel export format requested
|
|
873
|
-
check_requirements(
|
|
874
|
-
|
|
903
|
+
check_requirements(
|
|
904
|
+
["coremltools>=9.0", "numpy>=1.14.5,<=2.3.5"]
|
|
905
|
+
) # latest numpy 2.4.0rc1 breaks coremltools exports
|
|
906
|
+
import coremltools as ct
|
|
875
907
|
|
|
876
908
|
LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
|
|
877
909
|
assert not WINDOWS, "CoreML export is not supported on Windows, please run on macOS or Linux."
|
|
@@ -897,7 +929,7 @@ class Exporter:
|
|
|
897
929
|
model = IOSDetectModel(self.model, self.im, mlprogram=not mlmodel) if self.args.nms else self.model
|
|
898
930
|
else:
|
|
899
931
|
if self.args.nms:
|
|
900
|
-
LOGGER.warning(f"{prefix} 'nms=True' is only available for Detect models like '
|
|
932
|
+
LOGGER.warning(f"{prefix} 'nms=True' is only available for Detect models like 'yolo26n.pt'.")
|
|
901
933
|
# TODO CoreML Segment and Pose model pipelining
|
|
902
934
|
model = self.model
|
|
903
935
|
ts = torch.jit.trace(model.eval(), self.im, strict=False) # TorchScript model
|
|
@@ -917,7 +949,7 @@ class Exporter:
|
|
|
917
949
|
|
|
918
950
|
# Based on apple's documentation it is better to leave out the minimum_deployment target and let that get set
|
|
919
951
|
# Internally based on the model conversion and output type.
|
|
920
|
-
# Setting
|
|
952
|
+
# Setting minimum_deployment_target >= iOS16 will require setting compute_precision=ct.precision.FLOAT32.
|
|
921
953
|
# iOS16 adds in better support for FP16, but none of the CoreML NMS specifications handle FP16 as input.
|
|
922
954
|
ct_model = ct.convert(
|
|
923
955
|
ts,
|
|
@@ -967,12 +999,12 @@ class Exporter:
|
|
|
967
999
|
f_onnx = self.export_onnx() # run before TRT import https://github.com/ultralytics/ultralytics/issues/7016
|
|
968
1000
|
|
|
969
1001
|
try:
|
|
970
|
-
import tensorrt as trt
|
|
1002
|
+
import tensorrt as trt
|
|
971
1003
|
except ImportError:
|
|
972
1004
|
if LINUX:
|
|
973
1005
|
cuda_version = torch.version.cuda.split(".")[0]
|
|
974
1006
|
check_requirements(f"tensorrt-cu{cuda_version}>7.0.0,!=10.1.0")
|
|
975
|
-
import tensorrt as trt
|
|
1007
|
+
import tensorrt as trt
|
|
976
1008
|
check_version(trt.__version__, ">=7.0.0", hard=True)
|
|
977
1009
|
check_version(trt.__version__, "!=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
|
|
978
1010
|
|
|
@@ -1002,17 +1034,17 @@ class Exporter:
|
|
|
1002
1034
|
"""Export YOLO model to TensorFlow SavedModel format."""
|
|
1003
1035
|
cuda = torch.cuda.is_available()
|
|
1004
1036
|
try:
|
|
1005
|
-
import tensorflow as tf
|
|
1037
|
+
import tensorflow as tf
|
|
1006
1038
|
except ImportError:
|
|
1007
1039
|
check_requirements("tensorflow>=2.0.0,<=2.19.0")
|
|
1008
|
-
import tensorflow as tf
|
|
1040
|
+
import tensorflow as tf
|
|
1009
1041
|
check_requirements(
|
|
1010
1042
|
(
|
|
1011
1043
|
"tf_keras<=2.19.0", # required by 'onnx2tf' package
|
|
1012
1044
|
"sng4onnx>=1.0.1", # required by 'onnx2tf' package
|
|
1013
1045
|
"onnx_graphsurgeon>=0.3.26", # required by 'onnx2tf' package
|
|
1014
1046
|
"ai-edge-litert>=1.2.0" + (",<1.4.0" if MACOS else ""), # required by 'onnx2tf' package
|
|
1015
|
-
"onnx>=1.12.0",
|
|
1047
|
+
"onnx>=1.12.0,<2.0.0",
|
|
1016
1048
|
"onnx2tf>=1.26.3",
|
|
1017
1049
|
"onnxslim>=0.1.71",
|
|
1018
1050
|
"onnxruntime-gpu" if cuda else "onnxruntime",
|
|
@@ -1033,82 +1065,50 @@ class Exporter:
|
|
|
1033
1065
|
if f.is_dir():
|
|
1034
1066
|
shutil.rmtree(f) # delete output folder
|
|
1035
1067
|
|
|
1036
|
-
#
|
|
1037
|
-
|
|
1038
|
-
if
|
|
1039
|
-
|
|
1068
|
+
# Export to TF
|
|
1069
|
+
images = None
|
|
1070
|
+
if self.args.int8 and self.args.data:
|
|
1071
|
+
images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
|
|
1072
|
+
images = (
|
|
1073
|
+
torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz)
|
|
1074
|
+
.permute(0, 2, 3, 1)
|
|
1075
|
+
.numpy()
|
|
1076
|
+
.astype(np.float32)
|
|
1077
|
+
)
|
|
1040
1078
|
|
|
1041
1079
|
# Export to ONNX
|
|
1042
|
-
if
|
|
1080
|
+
if isinstance(self.model.model[-1], RTDETRDecoder):
|
|
1043
1081
|
self.args.opset = self.args.opset or 19
|
|
1044
1082
|
assert 16 <= self.args.opset <= 19, "RTDETR export requires opset>=16;<=19"
|
|
1045
1083
|
self.args.simplify = True
|
|
1046
|
-
f_onnx = self.export_onnx()
|
|
1047
|
-
|
|
1048
|
-
|
|
1049
|
-
|
|
1050
|
-
|
|
1051
|
-
|
|
1052
|
-
|
|
1053
|
-
|
|
1054
|
-
images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
|
|
1055
|
-
images = torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz).permute(
|
|
1056
|
-
0, 2, 3, 1
|
|
1057
|
-
)
|
|
1058
|
-
np.save(str(tmp_file), images.numpy().astype(np.float32)) # BHWC
|
|
1059
|
-
np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
|
|
1060
|
-
|
|
1061
|
-
import onnx2tf # scoped for after ONNX export for reduced conflict during import
|
|
1062
|
-
|
|
1063
|
-
LOGGER.info(f"{prefix} starting TFLite export with onnx2tf {onnx2tf.__version__}...")
|
|
1064
|
-
keras_model = onnx2tf.convert(
|
|
1065
|
-
input_onnx_file_path=f_onnx,
|
|
1066
|
-
output_folder_path=str(f),
|
|
1067
|
-
not_use_onnxsim=True,
|
|
1068
|
-
verbosity="error", # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
|
|
1069
|
-
output_integer_quantized_tflite=self.args.int8,
|
|
1070
|
-
custom_input_op_name_np_data_path=np_data,
|
|
1071
|
-
enable_batchmatmul_unfold=True and not self.args.int8, # fix lower no. of detected objects on GPU delegate
|
|
1072
|
-
output_signaturedefs=True, # fix error with Attention block group convolution
|
|
1073
|
-
disable_group_convolution=self.args.format in {"tfjs", "edgetpu"}, # fix error with group convolution
|
|
1084
|
+
f_onnx = self.export_onnx() # ensure ONNX is available
|
|
1085
|
+
keras_model = onnx2saved_model(
|
|
1086
|
+
f_onnx,
|
|
1087
|
+
f,
|
|
1088
|
+
int8=self.args.int8,
|
|
1089
|
+
images=images,
|
|
1090
|
+
disable_group_convolution=self.args.format in {"tfjs", "edgetpu"},
|
|
1091
|
+
prefix=prefix,
|
|
1074
1092
|
)
|
|
1075
1093
|
YAML.save(f / "metadata.yaml", self.metadata) # add metadata.yaml
|
|
1076
|
-
|
|
1077
|
-
# Remove/rename TFLite models
|
|
1078
|
-
if self.args.int8:
|
|
1079
|
-
tmp_file.unlink(missing_ok=True)
|
|
1080
|
-
for file in f.rglob("*_dynamic_range_quant.tflite"):
|
|
1081
|
-
file.rename(file.with_name(file.stem.replace("_dynamic_range_quant", "_int8") + file.suffix))
|
|
1082
|
-
for file in f.rglob("*_integer_quant_with_int16_act.tflite"):
|
|
1083
|
-
file.unlink() # delete extra fp16 activation TFLite files
|
|
1084
|
-
|
|
1085
1094
|
# Add TFLite metadata
|
|
1086
1095
|
for file in f.rglob("*.tflite"):
|
|
1087
|
-
|
|
1096
|
+
file.unlink() if "quant_with_int16_act.tflite" in str(file) else self._add_tflite_metadata(file)
|
|
1088
1097
|
|
|
1089
1098
|
return str(f), keras_model # or keras_model = tf.saved_model.load(f, tags=None, options=None)
|
|
1090
1099
|
|
|
1091
1100
|
@try_export
|
|
1092
1101
|
def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
|
|
1093
1102
|
"""Export YOLO model to TensorFlow GraphDef *.pb format https://github.com/leimao/Frozen-Graph-TensorFlow."""
|
|
1094
|
-
import tensorflow as tf # noqa
|
|
1095
|
-
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 # noqa
|
|
1096
|
-
|
|
1097
|
-
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
|
|
1098
1103
|
f = self.file.with_suffix(".pb")
|
|
1099
|
-
|
|
1100
|
-
m = tf.function(lambda x: keras_model(x)) # full model
|
|
1101
|
-
m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
|
|
1102
|
-
frozen_func = convert_variables_to_constants_v2(m)
|
|
1103
|
-
frozen_func.graph.as_graph_def()
|
|
1104
|
-
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
|
|
1104
|
+
keras2pb(keras_model, f, prefix)
|
|
1105
1105
|
return f
|
|
1106
1106
|
|
|
1107
1107
|
@try_export
|
|
1108
1108
|
def export_tflite(self, prefix=colorstr("TensorFlow Lite:")):
|
|
1109
1109
|
"""Export YOLO model to TensorFlow Lite format."""
|
|
1110
1110
|
# BUG https://github.com/ultralytics/ultralytics/issues/13436
|
|
1111
|
-
import tensorflow as tf
|
|
1111
|
+
import tensorflow as tf
|
|
1112
1112
|
|
|
1113
1113
|
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
|
|
1114
1114
|
saved_model = Path(str(self.file).replace(self.file.suffix, "_saved_model"))
|
|
@@ -1120,6 +1120,110 @@ class Exporter:
|
|
|
1120
1120
|
f = saved_model / f"{self.file.stem}_float32.tflite"
|
|
1121
1121
|
return str(f)
|
|
1122
1122
|
|
|
1123
|
+
@try_export
|
|
1124
|
+
def export_axelera(self, prefix=colorstr("Axelera:")):
|
|
1125
|
+
"""YOLO Axelera export."""
|
|
1126
|
+
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
|
|
1127
|
+
try:
|
|
1128
|
+
from axelera import compiler
|
|
1129
|
+
except ImportError:
|
|
1130
|
+
check_apt_requirements(
|
|
1131
|
+
["libllvm14", "libgirepository1.0-dev", "pkg-config", "libcairo2-dev", "build-essential", "cmake"]
|
|
1132
|
+
)
|
|
1133
|
+
|
|
1134
|
+
check_requirements(
|
|
1135
|
+
"axelera-voyager-sdk==1.5.2",
|
|
1136
|
+
cmds="--extra-index-url https://software.axelera.ai/artifactory/axelera-runtime-pypi "
|
|
1137
|
+
"--extra-index-url https://software.axelera.ai/artifactory/axelera-dev-pypi",
|
|
1138
|
+
)
|
|
1139
|
+
|
|
1140
|
+
from axelera import compiler
|
|
1141
|
+
from axelera.compiler import CompilerConfig
|
|
1142
|
+
|
|
1143
|
+
self.args.opset = 17 # hardcode opset for Axelera
|
|
1144
|
+
onnx_path = self.export_onnx()
|
|
1145
|
+
model_name = Path(onnx_path).stem
|
|
1146
|
+
export_path = Path(f"{model_name}_axelera_model")
|
|
1147
|
+
export_path.mkdir(exist_ok=True)
|
|
1148
|
+
|
|
1149
|
+
if "C2PSA" in self.model.__str__(): # YOLO11
|
|
1150
|
+
config = CompilerConfig(
|
|
1151
|
+
quantization_scheme="per_tensor_min_max",
|
|
1152
|
+
ignore_weight_buffers=False,
|
|
1153
|
+
resources_used=0.25,
|
|
1154
|
+
aipu_cores_used=1,
|
|
1155
|
+
multicore_mode="batch",
|
|
1156
|
+
output_axm_format=True,
|
|
1157
|
+
model_name=model_name,
|
|
1158
|
+
)
|
|
1159
|
+
else: # YOLOv8
|
|
1160
|
+
config = CompilerConfig(
|
|
1161
|
+
tiling_depth=6,
|
|
1162
|
+
split_buffer_promotion=True,
|
|
1163
|
+
resources_used=0.25,
|
|
1164
|
+
aipu_cores_used=1,
|
|
1165
|
+
multicore_mode="batch",
|
|
1166
|
+
output_axm_format=True,
|
|
1167
|
+
model_name=model_name,
|
|
1168
|
+
)
|
|
1169
|
+
|
|
1170
|
+
qmodel = compiler.quantize(
|
|
1171
|
+
model=onnx_path,
|
|
1172
|
+
calibration_dataset=self.get_int8_calibration_dataloader(prefix),
|
|
1173
|
+
config=config,
|
|
1174
|
+
transform_fn=self._transform_fn,
|
|
1175
|
+
)
|
|
1176
|
+
|
|
1177
|
+
compiler.compile(model=qmodel, config=config, output_dir=export_path)
|
|
1178
|
+
|
|
1179
|
+
axm_name = f"{model_name}.axm"
|
|
1180
|
+
axm_src = Path(axm_name)
|
|
1181
|
+
axm_dst = export_path / axm_name
|
|
1182
|
+
|
|
1183
|
+
if axm_src.exists():
|
|
1184
|
+
axm_src.replace(axm_dst)
|
|
1185
|
+
|
|
1186
|
+
YAML.save(export_path / "metadata.yaml", self.metadata)
|
|
1187
|
+
|
|
1188
|
+
return export_path
|
|
1189
|
+
|
|
1190
|
+
@try_export
|
|
1191
|
+
def export_executorch(self, prefix=colorstr("ExecuTorch:")):
|
|
1192
|
+
"""Exports a model to ExecuTorch (.pte) format into a dedicated directory and saves the required metadata,
|
|
1193
|
+
following Ultralytics conventions.
|
|
1194
|
+
"""
|
|
1195
|
+
LOGGER.info(f"\n{prefix} starting export with ExecuTorch...")
|
|
1196
|
+
assert TORCH_2_9, f"ExecuTorch export requires torch>=2.9.0 but torch=={TORCH_VERSION} is installed"
|
|
1197
|
+
|
|
1198
|
+
# BUG executorch build on arm64 Docker requires packaging>=22.0 https://github.com/pypa/setuptools/issues/4483
|
|
1199
|
+
if LINUX and ARM64 and IS_DOCKER:
|
|
1200
|
+
check_requirements("packaging>=22.0")
|
|
1201
|
+
|
|
1202
|
+
check_requirements("ruamel.yaml<0.19.0")
|
|
1203
|
+
check_requirements("executorch==1.0.1", "flatbuffers")
|
|
1204
|
+
# Pin numpy to avoid coremltools errors with numpy>=2.4.0, must be separate
|
|
1205
|
+
check_requirements("numpy<=2.3.5")
|
|
1206
|
+
|
|
1207
|
+
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
|
|
1208
|
+
from executorch.exir import to_edge_transform_and_lower
|
|
1209
|
+
|
|
1210
|
+
file_directory = Path(str(self.file).replace(self.file.suffix, "_executorch_model"))
|
|
1211
|
+
file_directory.mkdir(parents=True, exist_ok=True)
|
|
1212
|
+
|
|
1213
|
+
file_pte = file_directory / self.file.with_suffix(".pte").name
|
|
1214
|
+
sample_inputs = (self.im,)
|
|
1215
|
+
|
|
1216
|
+
et_program = to_edge_transform_and_lower(
|
|
1217
|
+
torch.export.export(self.model, sample_inputs), partitioner=[XnnpackPartitioner()]
|
|
1218
|
+
).to_executorch()
|
|
1219
|
+
|
|
1220
|
+
with open(file_pte, "wb") as file:
|
|
1221
|
+
file.write(et_program.buffer)
|
|
1222
|
+
|
|
1223
|
+
YAML.save(file_directory / "metadata.yaml", self.metadata)
|
|
1224
|
+
|
|
1225
|
+
return str(file_directory)
|
|
1226
|
+
|
|
1123
1227
|
@try_export
|
|
1124
1228
|
def export_edgetpu(self, tflite_model="", prefix=colorstr("Edge TPU:")):
|
|
1125
1229
|
"""Export YOLO model to Edge TPU format https://coral.ai/docs/edgetpu/models-intro/."""
|
|
@@ -1128,30 +1232,19 @@ class Exporter:
|
|
|
1128
1232
|
assert LINUX, f"export only supported on Linux. See {help_url}"
|
|
1129
1233
|
if subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True).returncode != 0:
|
|
1130
1234
|
LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
|
|
1235
|
+
sudo = "sudo " if is_sudo_available() else ""
|
|
1131
1236
|
for c in (
|
|
1132
|
-
"
|
|
1133
|
-
|
|
1134
|
-
"sudo
|
|
1135
|
-
"sudo apt-get update",
|
|
1136
|
-
"sudo apt-get install edgetpu-compiler",
|
|
1237
|
+
f"{sudo}mkdir -p /etc/apt/keyrings",
|
|
1238
|
+
f"curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | {sudo}gpg --dearmor -o /etc/apt/keyrings/google.gpg",
|
|
1239
|
+
f'echo "deb [signed-by=/etc/apt/keyrings/google.gpg] https://packages.cloud.google.com/apt coral-edgetpu-stable main" | {sudo}tee /etc/apt/sources.list.d/coral-edgetpu.list',
|
|
1137
1240
|
):
|
|
1138
|
-
subprocess.run(c
|
|
1139
|
-
|
|
1241
|
+
subprocess.run(c, shell=True, check=True)
|
|
1242
|
+
check_apt_requirements(["edgetpu-compiler"])
|
|
1140
1243
|
|
|
1244
|
+
ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
|
|
1141
1245
|
LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
|
|
1246
|
+
tflite2edgetpu(tflite_file=tflite_model, output_dir=tflite_model.parent, prefix=prefix)
|
|
1142
1247
|
f = str(tflite_model).replace(".tflite", "_edgetpu.tflite") # Edge TPU model
|
|
1143
|
-
|
|
1144
|
-
cmd = (
|
|
1145
|
-
"edgetpu_compiler "
|
|
1146
|
-
f'--out_dir "{Path(f).parent}" '
|
|
1147
|
-
"--show_operations "
|
|
1148
|
-
"--search_delegate "
|
|
1149
|
-
"--delegate_search_step 30 "
|
|
1150
|
-
"--timeout_sec 180 "
|
|
1151
|
-
f'"{tflite_model}"'
|
|
1152
|
-
)
|
|
1153
|
-
LOGGER.info(f"{prefix} running '{cmd}'")
|
|
1154
|
-
subprocess.run(cmd, shell=True)
|
|
1155
1248
|
self._add_tflite_metadata(f)
|
|
1156
1249
|
return f
|
|
1157
1250
|
|
|
@@ -1159,31 +1252,10 @@ class Exporter:
|
|
|
1159
1252
|
def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
|
|
1160
1253
|
"""Export YOLO model to TensorFlow.js format."""
|
|
1161
1254
|
check_requirements("tensorflowjs")
|
|
1162
|
-
import tensorflow as tf
|
|
1163
|
-
import tensorflowjs as tfjs # noqa
|
|
1164
1255
|
|
|
1165
|
-
LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
|
|
1166
1256
|
f = str(self.file).replace(self.file.suffix, "_web_model") # js dir
|
|
1167
1257
|
f_pb = str(self.file.with_suffix(".pb")) # *.pb path
|
|
1168
|
-
|
|
1169
|
-
gd = tf.Graph().as_graph_def() # TF GraphDef
|
|
1170
|
-
with open(f_pb, "rb") as file:
|
|
1171
|
-
gd.ParseFromString(file.read())
|
|
1172
|
-
outputs = ",".join(gd_outputs(gd))
|
|
1173
|
-
LOGGER.info(f"\n{prefix} output node names: {outputs}")
|
|
1174
|
-
|
|
1175
|
-
quantization = "--quantize_float16" if self.args.half else "--quantize_uint8" if self.args.int8 else ""
|
|
1176
|
-
with spaces_in_path(f_pb) as fpb_, spaces_in_path(f) as f_: # exporter can not handle spaces in path
|
|
1177
|
-
cmd = (
|
|
1178
|
-
"tensorflowjs_converter "
|
|
1179
|
-
f'--input_format=tf_frozen_model {quantization} --output_node_names={outputs} "{fpb_}" "{f_}"'
|
|
1180
|
-
)
|
|
1181
|
-
LOGGER.info(f"{prefix} running '{cmd}'")
|
|
1182
|
-
subprocess.run(cmd, shell=True)
|
|
1183
|
-
|
|
1184
|
-
if " " in f:
|
|
1185
|
-
LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{f}'.")
|
|
1186
|
-
|
|
1258
|
+
pb2tfjs(pb_file=f_pb, output_dir=f, half=self.args.half, int8=self.args.int8, prefix=prefix)
|
|
1187
1259
|
# Add metadata
|
|
1188
1260
|
YAML.save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
|
|
1189
1261
|
return f
|
|
@@ -1219,16 +1291,24 @@ class Exporter:
|
|
|
1219
1291
|
def export_imx(self, prefix=colorstr("IMX:")):
|
|
1220
1292
|
"""Export YOLO model to IMX format."""
|
|
1221
1293
|
assert LINUX, (
|
|
1222
|
-
"
|
|
1223
|
-
"See https://developer.aitrios.sony-semicon.com/en/
|
|
1294
|
+
"Export only supported on Linux."
|
|
1295
|
+
"See https://developer.aitrios.sony-semicon.com/en/docs/raspberry-pi-ai-camera/imx500-converter?version=3.17.3&progLang="
|
|
1224
1296
|
)
|
|
1297
|
+
assert not ARM64, "IMX export is not supported on ARM64 architectures."
|
|
1298
|
+
assert IS_PYTHON_MINIMUM_3_9, "IMX export is only supported on Python 3.9 or above."
|
|
1299
|
+
|
|
1225
1300
|
if getattr(self.model, "end2end", False):
|
|
1226
1301
|
raise ValueError("IMX export is not supported for end2end models.")
|
|
1227
1302
|
check_requirements(
|
|
1228
|
-
(
|
|
1303
|
+
(
|
|
1304
|
+
"model-compression-toolkit>=2.4.1",
|
|
1305
|
+
"edge-mdt-cl<1.1.0",
|
|
1306
|
+
"edge-mdt-tpc>=1.2.0",
|
|
1307
|
+
"pydantic<=2.11.7",
|
|
1308
|
+
)
|
|
1229
1309
|
)
|
|
1230
|
-
|
|
1231
|
-
check_requirements("
|
|
1310
|
+
|
|
1311
|
+
check_requirements("imx500-converter[pt]>=3.17.3")
|
|
1232
1312
|
|
|
1233
1313
|
# Install Java>=17
|
|
1234
1314
|
try:
|
|
@@ -1237,8 +1317,12 @@ class Exporter:
|
|
|
1237
1317
|
java_version = int(version_match.group(1)) if version_match else 0
|
|
1238
1318
|
assert java_version >= 17, "Java version too old"
|
|
1239
1319
|
except (FileNotFoundError, subprocess.CalledProcessError, AssertionError):
|
|
1240
|
-
|
|
1241
|
-
|
|
1320
|
+
if IS_UBUNTU or IS_DEBIAN_TRIXIE:
|
|
1321
|
+
LOGGER.info(f"\n{prefix} installing Java 21 for Ubuntu...")
|
|
1322
|
+
check_apt_requirements(["openjdk-21-jre"])
|
|
1323
|
+
elif IS_RASPBERRYPI or IS_DEBIAN_BOOKWORM:
|
|
1324
|
+
LOGGER.info(f"\n{prefix} installing Java 17 for Raspberry Pi or Debian ...")
|
|
1325
|
+
check_apt_requirements(["openjdk-17-jre"])
|
|
1242
1326
|
|
|
1243
1327
|
return torch2imx(
|
|
1244
1328
|
self.model,
|
|
@@ -1260,7 +1344,7 @@ class Exporter:
|
|
|
1260
1344
|
|
|
1261
1345
|
def _pipeline_coreml(self, model, weights_dir=None, prefix=colorstr("CoreML Pipeline:")):
|
|
1262
1346
|
"""Create CoreML pipeline with NMS for YOLO detection models."""
|
|
1263
|
-
import coremltools as ct
|
|
1347
|
+
import coremltools as ct
|
|
1264
1348
|
|
|
1265
1349
|
LOGGER.info(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
|
|
1266
1350
|
|
|
@@ -1353,6 +1437,14 @@ class Exporter:
|
|
|
1353
1437
|
LOGGER.info(f"{prefix} pipeline success")
|
|
1354
1438
|
return model
|
|
1355
1439
|
|
|
1440
|
+
@staticmethod
|
|
1441
|
+
def _transform_fn(data_item) -> np.ndarray:
|
|
1442
|
+
"""The transformation function for Axelera/OpenVINO quantization preprocessing."""
|
|
1443
|
+
data_item: torch.Tensor = data_item["img"] if isinstance(data_item, dict) else data_item
|
|
1444
|
+
assert data_item.dtype == torch.uint8, "Input image must be uint8 for the quantization preprocessing"
|
|
1445
|
+
im = data_item.numpy().astype(np.float32) / 255.0 # uint8 to fp16/32 and 0 - 255 to 0.0 - 1.0
|
|
1446
|
+
return im[None] if im.ndim == 3 else im
|
|
1447
|
+
|
|
1356
1448
|
def add_callback(self, event: str, callback):
|
|
1357
1449
|
"""Append the given callback to the specified event."""
|
|
1358
1450
|
self.callbacks[event].append(callback)
|
|
@@ -1367,8 +1459,7 @@ class IOSDetectModel(torch.nn.Module):
|
|
|
1367
1459
|
"""Wrap an Ultralytics YOLO model for Apple iOS CoreML export."""
|
|
1368
1460
|
|
|
1369
1461
|
def __init__(self, model, im, mlprogram=True):
|
|
1370
|
-
"""
|
|
1371
|
-
Initialize the IOSDetectModel class with a YOLO model and example image.
|
|
1462
|
+
"""Initialize the IOSDetectModel class with a YOLO model and example image.
|
|
1372
1463
|
|
|
1373
1464
|
Args:
|
|
1374
1465
|
model (torch.nn.Module): The YOLO model to wrap.
|
|
@@ -1402,8 +1493,7 @@ class NMSModel(torch.nn.Module):
|
|
|
1402
1493
|
"""Model wrapper with embedded NMS for Detect, Segment, Pose and OBB."""
|
|
1403
1494
|
|
|
1404
1495
|
def __init__(self, model, args):
|
|
1405
|
-
"""
|
|
1406
|
-
Initialize the NMSModel.
|
|
1496
|
+
"""Initialize the NMSModel.
|
|
1407
1497
|
|
|
1408
1498
|
Args:
|
|
1409
1499
|
model (torch.nn.Module): The model to wrap with NMS postprocessing.
|
|
@@ -1416,15 +1506,14 @@ class NMSModel(torch.nn.Module):
|
|
|
1416
1506
|
self.is_tf = self.args.format in frozenset({"saved_model", "tflite", "tfjs"})
|
|
1417
1507
|
|
|
1418
1508
|
def forward(self, x):
|
|
1419
|
-
"""
|
|
1420
|
-
Perform inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
|
|
1509
|
+
"""Perform inference with NMS post-processing. Supports Detect, Segment, OBB and Pose.
|
|
1421
1510
|
|
|
1422
1511
|
Args:
|
|
1423
1512
|
x (torch.Tensor): The preprocessed tensor with shape (N, 3, H, W).
|
|
1424
1513
|
|
|
1425
1514
|
Returns:
|
|
1426
|
-
(torch.Tensor): List of detections, each an (N, max_det, 4 + 2 + extra_shape) Tensor where N is the
|
|
1427
|
-
|
|
1515
|
+
(torch.Tensor): List of detections, each an (N, max_det, 4 + 2 + extra_shape) Tensor where N is the number
|
|
1516
|
+
of detections after NMS.
|
|
1428
1517
|
"""
|
|
1429
1518
|
from functools import partial
|
|
1430
1519
|
|
|
@@ -1455,17 +1544,16 @@ class NMSModel(torch.nn.Module):
|
|
|
1455
1544
|
box, score, cls, extra = box[mask], score[mask], cls[mask], extra[mask]
|
|
1456
1545
|
nmsbox = box.clone()
|
|
1457
1546
|
# `8` is the minimum value experimented to get correct NMS results for obb
|
|
1458
|
-
multiplier = 8 if self.obb else 1
|
|
1547
|
+
multiplier = 8 if self.obb else 1 / max(len(self.model.names), 1)
|
|
1459
1548
|
# Normalize boxes for NMS since large values for class offset causes issue with int8 quantization
|
|
1460
1549
|
if self.args.format == "tflite": # TFLite is already normalized
|
|
1461
1550
|
nmsbox *= multiplier
|
|
1462
1551
|
else:
|
|
1463
|
-
nmsbox = multiplier * nmsbox / torch.tensor(x.shape[2:], **kwargs).max()
|
|
1464
|
-
if not self.args.agnostic_nms: # class-
|
|
1552
|
+
nmsbox = multiplier * (nmsbox / torch.tensor(x.shape[2:], **kwargs).max())
|
|
1553
|
+
if not self.args.agnostic_nms: # class-wise NMS
|
|
1465
1554
|
end = 2 if self.obb else 4
|
|
1466
1555
|
# fully explicit expansion otherwise reshape error
|
|
1467
|
-
|
|
1468
|
-
cls_offset = cls.reshape(-1, 1).expand(nmsbox.shape[0], end)
|
|
1556
|
+
cls_offset = cls.view(cls.shape[0], 1).expand(cls.shape[0], end)
|
|
1469
1557
|
offbox = nmsbox[:, :end] + cls_offset * multiplier
|
|
1470
1558
|
nmsbox = torch.cat((offbox, nmsbox[:, end:]), dim=-1)
|
|
1471
1559
|
nms_fn = (
|