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.
- tests/__init__.py +22 -0
- tests/conftest.py +83 -0
- tests/test_cli.py +122 -0
- tests/test_cuda.py +155 -0
- tests/test_engine.py +131 -0
- tests/test_exports.py +216 -0
- tests/test_integrations.py +150 -0
- tests/test_python.py +615 -0
- tests/test_solutions.py +94 -0
- ultralytics/__init__.py +11 -8
- ultralytics/cfg/__init__.py +569 -131
- ultralytics/cfg/datasets/Argoverse.yaml +2 -1
- ultralytics/cfg/datasets/DOTAv1.5.yaml +3 -2
- ultralytics/cfg/datasets/DOTAv1.yaml +3 -2
- ultralytics/cfg/datasets/GlobalWheat2020.yaml +3 -2
- ultralytics/cfg/datasets/ImageNet.yaml +2 -1
- ultralytics/cfg/datasets/Objects365.yaml +5 -4
- ultralytics/cfg/datasets/SKU-110K.yaml +2 -1
- ultralytics/cfg/datasets/VOC.yaml +3 -2
- ultralytics/cfg/datasets/VisDrone.yaml +6 -5
- ultralytics/cfg/datasets/african-wildlife.yaml +25 -0
- ultralytics/cfg/datasets/brain-tumor.yaml +23 -0
- ultralytics/cfg/datasets/carparts-seg.yaml +3 -2
- ultralytics/cfg/datasets/coco-pose.yaml +7 -6
- ultralytics/cfg/datasets/coco.yaml +3 -2
- ultralytics/cfg/datasets/coco128-seg.yaml +4 -3
- ultralytics/cfg/datasets/coco128.yaml +4 -3
- ultralytics/cfg/datasets/coco8-pose.yaml +3 -2
- ultralytics/cfg/datasets/coco8-seg.yaml +3 -2
- ultralytics/cfg/datasets/coco8.yaml +3 -2
- ultralytics/cfg/datasets/crack-seg.yaml +3 -2
- ultralytics/cfg/datasets/dog-pose.yaml +24 -0
- ultralytics/cfg/datasets/dota8.yaml +3 -2
- ultralytics/cfg/datasets/hand-keypoints.yaml +26 -0
- ultralytics/cfg/datasets/lvis.yaml +1236 -0
- ultralytics/cfg/datasets/medical-pills.yaml +22 -0
- ultralytics/cfg/datasets/open-images-v7.yaml +2 -1
- ultralytics/cfg/datasets/package-seg.yaml +5 -4
- ultralytics/cfg/datasets/signature.yaml +21 -0
- ultralytics/cfg/datasets/tiger-pose.yaml +3 -2
- ultralytics/cfg/datasets/xView.yaml +2 -1
- ultralytics/cfg/default.yaml +14 -11
- ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +24 -0
- ultralytics/cfg/models/11/yolo11-cls.yaml +33 -0
- ultralytics/cfg/models/11/yolo11-obb.yaml +50 -0
- ultralytics/cfg/models/11/yolo11-pose.yaml +51 -0
- ultralytics/cfg/models/11/yolo11-seg.yaml +50 -0
- ultralytics/cfg/models/11/yolo11.yaml +50 -0
- ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +5 -2
- ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +5 -2
- ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +5 -2
- ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +5 -2
- ultralytics/cfg/models/v10/yolov10b.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10l.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10m.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10n.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10s.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10x.yaml +45 -0
- ultralytics/cfg/models/v3/yolov3-spp.yaml +5 -2
- ultralytics/cfg/models/v3/yolov3-tiny.yaml +5 -2
- ultralytics/cfg/models/v3/yolov3.yaml +5 -2
- ultralytics/cfg/models/v5/yolov5-p6.yaml +5 -2
- ultralytics/cfg/models/v5/yolov5.yaml +5 -2
- ultralytics/cfg/models/v6/yolov6.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-cls.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +6 -2
- ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +6 -2
- ultralytics/cfg/models/v8/yolov8-ghost.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-obb.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-p2.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-p6.yaml +10 -7
- ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-pose.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-seg.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-world.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-worldv2.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8.yaml +5 -2
- ultralytics/cfg/models/v9/yolov9c-seg.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9c.yaml +30 -25
- ultralytics/cfg/models/v9/yolov9e-seg.yaml +64 -0
- ultralytics/cfg/models/v9/yolov9e.yaml +46 -42
- ultralytics/cfg/models/v9/yolov9m.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9s.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9t.yaml +41 -0
- ultralytics/cfg/solutions/default.yaml +24 -0
- ultralytics/cfg/trackers/botsort.yaml +8 -5
- ultralytics/cfg/trackers/bytetrack.yaml +8 -5
- ultralytics/data/__init__.py +14 -3
- ultralytics/data/annotator.py +37 -15
- ultralytics/data/augment.py +1783 -289
- ultralytics/data/base.py +62 -27
- ultralytics/data/build.py +37 -8
- ultralytics/data/converter.py +196 -36
- ultralytics/data/dataset.py +233 -94
- ultralytics/data/loaders.py +199 -96
- ultralytics/data/split_dota.py +39 -29
- ultralytics/data/utils.py +111 -41
- ultralytics/engine/__init__.py +1 -1
- ultralytics/engine/exporter.py +579 -244
- ultralytics/engine/model.py +604 -252
- ultralytics/engine/predictor.py +22 -11
- ultralytics/engine/results.py +1228 -218
- ultralytics/engine/trainer.py +191 -129
- ultralytics/engine/tuner.py +18 -18
- ultralytics/engine/validator.py +18 -15
- ultralytics/hub/__init__.py +31 -13
- ultralytics/hub/auth.py +11 -7
- ultralytics/hub/google/__init__.py +159 -0
- ultralytics/hub/session.py +128 -94
- ultralytics/hub/utils.py +20 -21
- ultralytics/models/__init__.py +4 -2
- ultralytics/models/fastsam/__init__.py +2 -3
- ultralytics/models/fastsam/model.py +26 -4
- ultralytics/models/fastsam/predict.py +127 -63
- ultralytics/models/fastsam/utils.py +1 -44
- ultralytics/models/fastsam/val.py +1 -1
- ultralytics/models/nas/__init__.py +1 -1
- ultralytics/models/nas/model.py +21 -10
- ultralytics/models/nas/predict.py +3 -6
- ultralytics/models/nas/val.py +4 -4
- ultralytics/models/rtdetr/__init__.py +1 -1
- ultralytics/models/rtdetr/model.py +1 -1
- ultralytics/models/rtdetr/predict.py +6 -8
- ultralytics/models/rtdetr/train.py +6 -2
- ultralytics/models/rtdetr/val.py +3 -3
- ultralytics/models/sam/__init__.py +3 -3
- ultralytics/models/sam/amg.py +29 -23
- ultralytics/models/sam/build.py +211 -13
- ultralytics/models/sam/model.py +91 -30
- ultralytics/models/sam/modules/__init__.py +1 -1
- ultralytics/models/sam/modules/blocks.py +1129 -0
- ultralytics/models/sam/modules/decoders.py +381 -53
- ultralytics/models/sam/modules/encoders.py +515 -324
- ultralytics/models/sam/modules/memory_attention.py +237 -0
- ultralytics/models/sam/modules/sam.py +969 -21
- ultralytics/models/sam/modules/tiny_encoder.py +425 -154
- ultralytics/models/sam/modules/transformer.py +159 -60
- ultralytics/models/sam/modules/utils.py +293 -0
- ultralytics/models/sam/predict.py +1263 -132
- ultralytics/models/utils/__init__.py +1 -1
- ultralytics/models/utils/loss.py +36 -24
- ultralytics/models/utils/ops.py +3 -7
- ultralytics/models/yolo/__init__.py +3 -3
- ultralytics/models/yolo/classify/__init__.py +1 -1
- ultralytics/models/yolo/classify/predict.py +7 -8
- ultralytics/models/yolo/classify/train.py +17 -22
- ultralytics/models/yolo/classify/val.py +8 -4
- ultralytics/models/yolo/detect/__init__.py +1 -1
- ultralytics/models/yolo/detect/predict.py +3 -5
- ultralytics/models/yolo/detect/train.py +11 -4
- ultralytics/models/yolo/detect/val.py +90 -52
- ultralytics/models/yolo/model.py +14 -9
- ultralytics/models/yolo/obb/__init__.py +1 -1
- ultralytics/models/yolo/obb/predict.py +2 -2
- ultralytics/models/yolo/obb/train.py +5 -3
- ultralytics/models/yolo/obb/val.py +41 -23
- ultralytics/models/yolo/pose/__init__.py +1 -1
- ultralytics/models/yolo/pose/predict.py +3 -5
- ultralytics/models/yolo/pose/train.py +2 -2
- ultralytics/models/yolo/pose/val.py +51 -17
- ultralytics/models/yolo/segment/__init__.py +1 -1
- ultralytics/models/yolo/segment/predict.py +3 -5
- ultralytics/models/yolo/segment/train.py +2 -2
- ultralytics/models/yolo/segment/val.py +60 -19
- ultralytics/models/yolo/world/__init__.py +5 -0
- ultralytics/models/yolo/world/train.py +92 -0
- ultralytics/models/yolo/world/train_world.py +109 -0
- ultralytics/nn/__init__.py +1 -1
- ultralytics/nn/autobackend.py +228 -93
- ultralytics/nn/modules/__init__.py +39 -14
- ultralytics/nn/modules/activation.py +21 -0
- ultralytics/nn/modules/block.py +526 -66
- ultralytics/nn/modules/conv.py +24 -7
- ultralytics/nn/modules/head.py +177 -34
- ultralytics/nn/modules/transformer.py +6 -5
- ultralytics/nn/modules/utils.py +1 -2
- ultralytics/nn/tasks.py +226 -82
- ultralytics/solutions/__init__.py +30 -1
- ultralytics/solutions/ai_gym.py +96 -143
- ultralytics/solutions/analytics.py +247 -0
- ultralytics/solutions/distance_calculation.py +78 -135
- ultralytics/solutions/heatmap.py +93 -247
- ultralytics/solutions/object_counter.py +184 -259
- ultralytics/solutions/parking_management.py +246 -0
- ultralytics/solutions/queue_management.py +112 -0
- ultralytics/solutions/region_counter.py +116 -0
- ultralytics/solutions/security_alarm.py +144 -0
- ultralytics/solutions/solutions.py +178 -0
- ultralytics/solutions/speed_estimation.py +86 -174
- ultralytics/solutions/streamlit_inference.py +190 -0
- ultralytics/solutions/trackzone.py +68 -0
- ultralytics/trackers/__init__.py +1 -1
- ultralytics/trackers/basetrack.py +32 -13
- ultralytics/trackers/bot_sort.py +61 -28
- ultralytics/trackers/byte_tracker.py +83 -51
- ultralytics/trackers/track.py +21 -6
- ultralytics/trackers/utils/__init__.py +1 -1
- ultralytics/trackers/utils/gmc.py +62 -48
- ultralytics/trackers/utils/kalman_filter.py +166 -35
- ultralytics/trackers/utils/matching.py +40 -21
- ultralytics/utils/__init__.py +511 -239
- ultralytics/utils/autobatch.py +40 -22
- ultralytics/utils/benchmarks.py +266 -85
- ultralytics/utils/callbacks/__init__.py +1 -1
- ultralytics/utils/callbacks/base.py +1 -3
- ultralytics/utils/callbacks/clearml.py +7 -6
- ultralytics/utils/callbacks/comet.py +39 -17
- ultralytics/utils/callbacks/dvc.py +1 -1
- ultralytics/utils/callbacks/hub.py +16 -16
- ultralytics/utils/callbacks/mlflow.py +28 -24
- ultralytics/utils/callbacks/neptune.py +6 -2
- ultralytics/utils/callbacks/raytune.py +3 -4
- ultralytics/utils/callbacks/tensorboard.py +18 -18
- ultralytics/utils/callbacks/wb.py +27 -20
- ultralytics/utils/checks.py +172 -100
- ultralytics/utils/dist.py +2 -1
- ultralytics/utils/downloads.py +40 -34
- ultralytics/utils/errors.py +1 -1
- ultralytics/utils/files.py +72 -38
- ultralytics/utils/instance.py +41 -19
- ultralytics/utils/loss.py +83 -55
- ultralytics/utils/metrics.py +61 -56
- ultralytics/utils/ops.py +94 -89
- ultralytics/utils/patches.py +30 -14
- ultralytics/utils/plotting.py +600 -269
- ultralytics/utils/tal.py +67 -26
- ultralytics/utils/torch_utils.py +305 -112
- ultralytics/utils/triton.py +2 -1
- ultralytics/utils/tuner.py +21 -12
- ultralytics-8.3.63.dist-info/METADATA +370 -0
- ultralytics-8.3.63.dist-info/RECORD +241 -0
- {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/WHEEL +1 -1
- ultralytics/data/explorer/__init__.py +0 -5
- ultralytics/data/explorer/explorer.py +0 -472
- ultralytics/data/explorer/gui/__init__.py +0 -1
- ultralytics/data/explorer/gui/dash.py +0 -268
- ultralytics/data/explorer/utils.py +0 -166
- ultralytics/models/fastsam/prompt.py +0 -357
- ultralytics-8.1.29.dist-info/METADATA +0 -373
- ultralytics-8.1.29.dist-info/RECORD +0 -197
- {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/LICENSE +0 -0
- {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/entry_points.txt +0 -0
- {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/top_level.txt +0 -0
ultralytics/engine/exporter.py
CHANGED
@@ -1,55 +1,60 @@
|
|
1
|
-
# Ultralytics
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
"""
|
3
|
-
Export a
|
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 | - |
|
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
|
-
|
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('
|
28
|
+
model = YOLO('yolo11n.pt')
|
27
29
|
results = model.export(format='onnx')
|
28
30
|
|
29
31
|
CLI:
|
30
|
-
$ yolo mode=export model=
|
32
|
+
$ yolo mode=export model=yolo11n.pt format=onnx
|
31
33
|
|
32
34
|
Inference:
|
33
|
-
$ yolo predict model=
|
34
|
-
|
35
|
-
|
36
|
-
|
37
|
-
|
38
|
-
|
39
|
-
|
40
|
-
|
41
|
-
|
42
|
-
|
43
|
-
|
44
|
-
|
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 ../../
|
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
|
87
|
-
|
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
|
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
|
-
"""
|
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
|
-
["
|
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
|
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
|
-
"""
|
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.
|
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
|
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
|
-
|
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
|
211
|
+
if fmt in {"tensorrt", "trt"}: # 'engine' aliases
|
175
212
|
fmt = "engine"
|
176
|
-
if fmt in
|
213
|
+
if fmt in {"mlmodel", "mlpackage", "mlprogram", "apple", "ios", "coreml"}: # 'coreml' aliases
|
177
214
|
fmt = "coreml"
|
178
|
-
|
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
|
-
|
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
|
-
#
|
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
|
203
|
-
|
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
|
-
|
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
|
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
|
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
|
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[
|
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
|
436
|
+
f"\nExport complete ({time.time() - t:.1f}s)"
|
324
437
|
f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
|
325
|
-
f
|
326
|
-
f
|
327
|
-
f
|
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
|
-
"""
|
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
|
-
"""
|
488
|
+
"""YOLO ONNX export."""
|
353
489
|
requirements = ["onnx>=1.12.0"]
|
354
490
|
if self.args.simplify:
|
355
|
-
requirements += ["
|
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
|
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.
|
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
|
-
"""
|
414
|
-
check_requirements("openvino>=2024.
|
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
|
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("
|
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
|
-
|
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
|
-
|
454
|
-
|
455
|
-
|
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,
|
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
|
-
"""
|
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
|
-
|
532
|
-
|
533
|
-
|
534
|
-
|
535
|
-
|
536
|
-
|
537
|
-
|
538
|
-
|
539
|
-
|
540
|
-
|
541
|
-
|
542
|
-
|
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
|
547
|
-
f
|
548
|
-
f
|
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
|
553
|
-
f
|
554
|
-
f
|
555
|
-
f
|
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
|
-
"""
|
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
|
-
"""
|
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
|
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("
|
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
|
-
|
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
|
-
|
679
|
-
|
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,
|
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
|
-
|
705
|
-
|
706
|
-
|
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
|
-
|
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
|
-
"""
|
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 = "
|
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
|
740
|
-
"
|
741
|
-
"
|
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
|
-
)
|
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
|
-
"
|
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
|
-
|
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 =
|
791
|
-
|
792
|
-
|
793
|
-
|
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=
|
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)
|
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
|
-
"""
|
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
|
-
"""
|
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
|
-
"""
|
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
|
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 =
|
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
|
-
"""
|
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
|
-
#
|
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
|
-
|
937
|
-
|
938
|
-
|
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 =
|
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 =
|
1303
|
+
label_file = schema.AssociatedFileT()
|
953
1304
|
label_file.name = tmp_file.name
|
954
|
-
label_file.type =
|
1305
|
+
label_file.type = schema.AssociatedFileType.TENSOR_AXIS_LABELS
|
955
1306
|
|
956
1307
|
# Create input info
|
957
|
-
input_meta =
|
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 =
|
961
|
-
input_meta.content.contentProperties =
|
962
|
-
input_meta.content.contentProperties.colorSpace =
|
963
|
-
input_meta.content.contentPropertiesType =
|
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 =
|
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 =
|
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 =
|
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),
|
1334
|
+
b.Finish(model_meta.Pack(b), metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
|
984
1335
|
metadata_buf = b.Output()
|
985
1336
|
|
986
|
-
populator =
|
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
|
-
"""
|
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)
|