ultralytics 8.3.25 Alibaba MNN export and predict support (#16802)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Francesco Mattioli <Francesco.mttl@gmail.com> Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com> Co-authored-by: Laughing-q <1185102784@qq.com> Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
王召德
GitHub
parent
11b4194344
commit
9c72d94ba4
14 changed files with 465 additions and 39 deletions
|
|
@ -1,6 +1,6 @@
|
|||
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
||||
|
||||
__version__ = "8.3.24"
|
||||
__version__ = "8.3.25"
|
||||
|
||||
import os
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ TensorFlow Lite | `tflite` | yolo11n.tflite
|
|||
TensorFlow Edge TPU | `edgetpu` | yolo11n_edgetpu.tflite
|
||||
TensorFlow.js | `tfjs` | yolo11n_web_model/
|
||||
PaddlePaddle | `paddle` | yolo11n_paddle_model/
|
||||
MNN | `mnn` | yolo11n.mnn
|
||||
NCNN | `ncnn` | yolo11n_ncnn_model/
|
||||
|
||||
Requirements:
|
||||
|
|
@ -41,6 +42,7 @@ Inference:
|
|||
yolo11n.tflite # TensorFlow Lite
|
||||
yolo11n_edgetpu.tflite # TensorFlow Edge TPU
|
||||
yolo11n_paddle_model # PaddlePaddle
|
||||
yolo11n.mnn # MNN
|
||||
yolo11n_ncnn_model # NCNN
|
||||
|
||||
TensorFlow.js:
|
||||
|
|
@ -109,6 +111,7 @@ def export_formats():
|
|||
["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", True, False],
|
||||
["TensorFlow.js", "tfjs", "_web_model", True, False],
|
||||
["PaddlePaddle", "paddle", "_paddle_model", True, True],
|
||||
["MNN", "mnn", ".mnn", True, True],
|
||||
["NCNN", "ncnn", "_ncnn_model", True, True],
|
||||
]
|
||||
return dict(zip(["Format", "Argument", "Suffix", "CPU", "GPU"], zip(*x)))
|
||||
|
|
@ -190,7 +193,9 @@ class Exporter:
|
|||
flags = [x == fmt for x in fmts]
|
||||
if sum(flags) != 1:
|
||||
raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
|
||||
jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, ncnn = flags # export booleans
|
||||
jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, mnn, ncnn = (
|
||||
flags # export booleans
|
||||
)
|
||||
is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))
|
||||
|
||||
# Device
|
||||
|
|
@ -333,8 +338,10 @@ class Exporter:
|
|||
f[9], _ = self.export_tfjs()
|
||||
if paddle: # PaddlePaddle
|
||||
f[10], _ = self.export_paddle()
|
||||
if mnn: # MNN
|
||||
f[11], _ = self.export_mnn()
|
||||
if ncnn: # NCNN
|
||||
f[11], _ = self.export_ncnn()
|
||||
f[12], _ = self.export_ncnn()
|
||||
|
||||
# Finish
|
||||
f = [str(x) for x in f if x] # filter out '' and None
|
||||
|
|
@ -541,6 +548,32 @@ class Exporter:
|
|||
yaml_save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
|
||||
return f, None
|
||||
|
||||
@try_export
|
||||
def export_mnn(self, prefix=colorstr("MNN:")):
|
||||
"""YOLOv8 MNN export using MNN https://github.com/alibaba/MNN."""
|
||||
f_onnx, _ = self.export_onnx() # get onnx model first
|
||||
|
||||
check_requirements("MNN>=2.9.6")
|
||||
import MNN # noqa
|
||||
from MNN.tools import mnnconvert
|
||||
|
||||
# Setup and checks
|
||||
LOGGER.info(f"\n{prefix} starting export with MNN {MNN.version()}...")
|
||||
assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
|
||||
f = str(self.file.with_suffix(".mnn")) # MNN model file
|
||||
args = ["", "-f", "ONNX", "--modelFile", f_onnx, "--MNNModel", f, "--bizCode", json.dumps(self.metadata)]
|
||||
if self.args.int8:
|
||||
args.append("--weightQuantBits")
|
||||
args.append("8")
|
||||
if self.args.half:
|
||||
args.append("--fp16")
|
||||
mnnconvert.convert(args)
|
||||
# remove scratch file for model convert optimize
|
||||
convert_scratch = Path(self.file.parent / ".__convert_external_data.bin")
|
||||
if convert_scratch.exists():
|
||||
convert_scratch.unlink()
|
||||
return f, None
|
||||
|
||||
@try_export
|
||||
def export_ncnn(self, prefix=colorstr("NCNN:")):
|
||||
"""YOLO NCNN export using PNNX https://github.com/pnnx/pnnx."""
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ Usage - formats:
|
|||
yolov8n.tflite # TensorFlow Lite
|
||||
yolov8n_edgetpu.tflite # TensorFlow Edge TPU
|
||||
yolov8n_paddle_model # PaddlePaddle
|
||||
yolov8n.mnn # MNN
|
||||
yolov8n_ncnn_model # NCNN
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ Usage - formats:
|
|||
yolov8n.tflite # TensorFlow Lite
|
||||
yolov8n_edgetpu.tflite # TensorFlow Edge TPU
|
||||
yolov8n_paddle_model # PaddlePaddle
|
||||
yolov8n.mnn # MNN
|
||||
yolov8n_ncnn_model # NCNN
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -59,21 +59,22 @@ class AutoBackend(nn.Module):
|
|||
range of formats, each with specific naming conventions as outlined below:
|
||||
|
||||
Supported Formats and Naming Conventions:
|
||||
| Format | File Suffix |
|
||||
|-----------------------|------------------|
|
||||
| PyTorch | *.pt |
|
||||
| TorchScript | *.torchscript |
|
||||
| ONNX Runtime | *.onnx |
|
||||
| ONNX OpenCV DNN | *.onnx (dnn=True)|
|
||||
| OpenVINO | *openvino_model/ |
|
||||
| CoreML | *.mlpackage |
|
||||
| TensorRT | *.engine |
|
||||
| TensorFlow SavedModel | *_saved_model |
|
||||
| TensorFlow GraphDef | *.pb |
|
||||
| TensorFlow Lite | *.tflite |
|
||||
| TensorFlow Edge TPU | *_edgetpu.tflite |
|
||||
| PaddlePaddle | *_paddle_model |
|
||||
| NCNN | *_ncnn_model |
|
||||
| Format | File Suffix |
|
||||
|-----------------------|-------------------|
|
||||
| PyTorch | *.pt |
|
||||
| TorchScript | *.torchscript |
|
||||
| ONNX Runtime | *.onnx |
|
||||
| ONNX OpenCV DNN | *.onnx (dnn=True) |
|
||||
| OpenVINO | *openvino_model/ |
|
||||
| CoreML | *.mlpackage |
|
||||
| TensorRT | *.engine |
|
||||
| TensorFlow SavedModel | *_saved_model/ |
|
||||
| TensorFlow GraphDef | *.pb |
|
||||
| TensorFlow Lite | *.tflite |
|
||||
| TensorFlow Edge TPU | *_edgetpu.tflite |
|
||||
| PaddlePaddle | *_paddle_model/ |
|
||||
| MNN | *.mnn |
|
||||
| NCNN | *_ncnn_model/ |
|
||||
|
||||
This class offers dynamic backend switching capabilities based on the input model format, making it easier to deploy
|
||||
models across various platforms.
|
||||
|
|
@ -120,6 +121,7 @@ class AutoBackend(nn.Module):
|
|||
edgetpu,
|
||||
tfjs,
|
||||
paddle,
|
||||
mnn,
|
||||
ncnn,
|
||||
triton,
|
||||
) = self._model_type(w)
|
||||
|
|
@ -403,6 +405,26 @@ class AutoBackend(nn.Module):
|
|||
output_names = predictor.get_output_names()
|
||||
metadata = w.parents[1] / "metadata.yaml"
|
||||
|
||||
# MNN
|
||||
elif mnn:
|
||||
LOGGER.info(f"Loading {w} for MNN inference...")
|
||||
check_requirements("MNN") # requires MNN
|
||||
import os
|
||||
|
||||
import MNN
|
||||
|
||||
config = {}
|
||||
config["precision"] = "low"
|
||||
config["backend"] = "CPU"
|
||||
config["numThread"] = (os.cpu_count() + 1) // 2
|
||||
rt = MNN.nn.create_runtime_manager((config,))
|
||||
net = MNN.nn.load_module_from_file(w, [], [], runtime_manager=rt, rearrange=True)
|
||||
|
||||
def torch_to_mnn(x):
|
||||
return MNN.expr.const(x.data_ptr(), x.shape)
|
||||
|
||||
metadata = json.loads(net.get_info()["bizCode"])
|
||||
|
||||
# NCNN
|
||||
elif ncnn:
|
||||
LOGGER.info(f"Loading {w} for NCNN inference...")
|
||||
|
|
@ -590,6 +612,12 @@ class AutoBackend(nn.Module):
|
|||
self.predictor.run()
|
||||
y = [self.predictor.get_output_handle(x).copy_to_cpu() for x in self.output_names]
|
||||
|
||||
# MNN
|
||||
elif self.mnn:
|
||||
input_var = self.torch_to_mnn(im)
|
||||
output_var = self.net.onForward([input_var])
|
||||
y = [x.read() for x in output_var]
|
||||
|
||||
# NCNN
|
||||
elif self.ncnn:
|
||||
mat_in = self.pyncnn.Mat(im[0].cpu().numpy())
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ TensorFlow Lite | `tflite` | yolov8n.tflite
|
|||
TensorFlow Edge TPU | `edgetpu` | yolov8n_edgetpu.tflite
|
||||
TensorFlow.js | `tfjs` | yolov8n_web_model/
|
||||
PaddlePaddle | `paddle` | yolov8n_paddle_model/
|
||||
MNN | `mnn` | yolov8n.mnn
|
||||
NCNN | `ncnn` | yolov8n_ncnn_model/
|
||||
"""
|
||||
|
||||
|
|
@ -111,8 +112,8 @@ def benchmark(
|
|||
assert not isinstance(model, YOLOWorld), "YOLOWorldv2 Paddle exports not supported yet"
|
||||
assert not is_end2end, "End-to-end models not supported by PaddlePaddle yet"
|
||||
assert LINUX or MACOS, "Windows Paddle exports not supported yet"
|
||||
if i in {12}: # NCNN
|
||||
assert not isinstance(model, YOLOWorld), "YOLOWorldv2 NCNN exports not supported yet"
|
||||
if i in {12, 13}: # MNN, NCNN
|
||||
assert not isinstance(model, YOLOWorld), "YOLOWorldv2 MNN, NCNN exports not supported yet"
|
||||
if "cpu" in device.type:
|
||||
assert cpu, "inference not supported on CPU"
|
||||
if "cuda" in device.type:
|
||||
|
|
@ -132,7 +133,7 @@ def benchmark(
|
|||
assert model.task != "pose" or i != 7, "GraphDef Pose inference is not supported"
|
||||
assert i not in {9, 10}, "inference not supported" # Edge TPU and TF.js are unsupported
|
||||
assert i != 5 or platform.system() == "Darwin", "inference only supported on macOS>=10.13" # CoreML
|
||||
if i in {12}:
|
||||
if i in {13}:
|
||||
assert not is_end2end, "End-to-end torch.topk operation is not supported for NCNN prediction yet"
|
||||
exported_model.predict(ASSETS / "bus.jpg", imgsz=imgsz, device=device, half=half)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue