Elgato_dark/ultralytics
1
0
Fork 0
mirror of https://github.com/ultralytics/ultralytics synced 2026-04-21 14:07:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

Refactor exports for saved_model/pb/edgetpu/tfjs formats (#22115)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.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:
Jing Qiu 2025-10-30 22:45:58 +08:00 committed by GitHub
parent 38575453c6
commit 82d3b50996
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 576 additions and 375 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
docs/en/reference/engine/exporter.md
View file

@ -35,10 +35,6 @@ keywords: YOLOv8, export formats, ONNX, TensorRT, CoreML, machine learning model
<br><br><hr><br>
## ::: ultralytics.engine.exporter.gd_outputs
<br><br><hr><br>
## ::: ultralytics.engine.exporter.try_export
<br><br>

20
docs/en/reference/utils/export/engine.md Normal file
View file

@ -0,0 +1,20 @@
---
description: TensorRT engine export utilities for converting ONNX models to optimized TensorRT engines. Provides functions for ONNX export from PyTorch models and TensorRT engine generation with support for FP16/INT8 quantization, dynamic shapes, DLA acceleration, and INT8 calibration for NVIDIA GPU inference optimization.
keywords: Ultralytics, TensorRT export, ONNX export, PyTorch to ONNX, quantization, FP16, INT8, dynamic shapes, DLA acceleration, GPU inference, model optimization, calibration, NVIDIA, inference engine, model export
---
# Reference for `ultralytics/utils/export/engine.py`
!!! note
This file is available at [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/engine.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/engine.py). If you spot a problem please help fix it by [contributing](https://docs.ultralytics.com/help/contributing/) a [Pull Request](https://github.com/ultralytics/ultralytics/edit/main/ultralytics/utils/export/engine.py) 🛠️. Thank you 🙏!
<br>
## ::: ultralytics.utils.export.engine.torch2onnx
<br><br><hr><br>
## ::: ultralytics.utils.export.engine.onnx2engine
<br><br>

44
docs/en/reference/utils/export/tensorflow.md Normal file
View file

@ -0,0 +1,44 @@
---
description: TensorFlow export utilities for converting PyTorch models to various TensorFlow formats. Provides functions for converting models to TensorFlow SavedModel, Protocol Buffer (.pb), TensorFlow Lite, Edge TPU, and TensorFlow.js formats via ONNX intermediate representation with support for INT8 quantization and calibration.
keywords: Ultralytics, TensorFlow, SavedModel, Protocol Buffer, TensorFlow Lite, TFLite, Edge TPU, TensorFlow.js, ONNX conversion, PyTorch to TensorFlow, INT8 quantization, model calibration, frozen graph, onnx2tf, model export, web deployment, mobile deployment
---
# Reference for `ultralytics/utils/export/tensorflow.py`
!!! note
This file is available at [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/tensorflow.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/export/tensorflow.py). If you spot a problem please help fix it by [contributing](https://docs.ultralytics.com/help/contributing/) a [Pull Request](https://github.com/ultralytics/ultralytics/edit/main/ultralytics/utils/export/tensorflow.py) 🛠️. Thank you 🙏!
<br>
## ::: ultralytics.utils.export.tensorflow.tf_wrapper
<br><br><hr><br>
## ::: ultralytics.utils.export.tensorflow._tf_inference
<br><br><hr><br>
## ::: ultralytics.utils.export.tensorflow.tf_kpts_decode
<br><br><hr><br>
## ::: ultralytics.utils.export.tensorflow.onnx2saved_model
<br><br><hr><br>
## ::: ultralytics.utils.export.tensorflow.keras2pb
<br><br><hr><br>
## ::: ultralytics.utils.export.tensorflow.tflite2edgetpu
<br><br><hr><br>
## ::: ultralytics.utils.export.tensorflow.pb2tfjs
<br><br><hr><br>
## ::: ultralytics.utils.export.tensorflow.gd_outputs
<br><br>

3
mkdocs.yml
View file

@ -693,8 +693,9 @@ nav:
- errors: reference/utils/errors.md
- events: reference/utils/events.md
- export:
- __init__: reference/utils/export/__init__.md
- engine: reference/utils/export/engine.md
- imx: reference/utils/export/imx.md
- tensorflow: reference/utils/export/tensorflow.md
- files: reference/utils/files.md
- git: reference/utils/git.md
- instance: reference/utils/instance.md

139
ultralytics/engine/exporter.py
View file

@ -107,9 +107,17 @@ from ultralytics.utils.checks import (
is_intel,
is_sudo_available,
)
from ultralytics.utils.downloads import attempt_download_asset, get_github_assets, safe_download
from ultralytics.utils.export import onnx2engine, torch2imx, torch2onnx
from ultralytics.utils.files import file_size, spaces_in_path
from ultralytics.utils.downloads import get_github_assets, safe_download
from ultralytics.utils.export import (
keras2pb,
onnx2engine,
onnx2saved_model,
pb2tfjs,
tflite2edgetpu,
torch2imx,
torch2onnx,
)
from ultralytics.utils.files import file_size
from ultralytics.utils.metrics import batch_probiou
from ultralytics.utils.nms import TorchNMS
from ultralytics.utils.ops import Profile
@ -206,15 +214,6 @@ def validate_args(format, passed_args, valid_args):
assert arg in valid_args, f"ERROR ❌️ argument '{arg}' is not supported for format='{format}'"
def gd_outputs(gd):
"""Return TensorFlow GraphDef model output node names."""
name_list, input_list = [], []
for node in gd.node: # tensorflow.core.framework.node_def_pb2.NodeDef
name_list.append(node.name)
input_list.extend(node.input)
return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))
def try_export(inner_func):
"""YOLO export decorator, i.e. @try_export."""
inner_args = get_default_args(inner_func)
@ -461,6 +460,10 @@ class Exporter:
from ultralytics.utils.export.imx import FXModel
model = FXModel(model, self.imgsz)
if tflite or edgetpu:
from ultralytics.utils.export.tensorflow import tf_wrapper
model = tf_wrapper(model)
for m in model.modules():
if isinstance(m, Classify):
m.export = True
@ -642,7 +645,7 @@ class Exporter:
assert TORCH_1_13, f"'nms=True' ONNX export requires torch>=1.13 (found torch=={TORCH_VERSION})"
f = str(self.file.with_suffix(".onnx"))
output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
output_names = ["output0", "output1"] if self.model.task == "segment" else ["output0"]
dynamic = self.args.dynamic
if dynamic:
dynamic = {"images": {0: "batch", 2: "height", 3: "width"}} # shape(1,3,640,640)
@ -1053,75 +1056,43 @@ class Exporter:
if f.is_dir():
shutil.rmtree(f) # delete output folder
# Pre-download calibration file to fix https://github.com/PINTO0309/onnx2tf/issues/545
onnx2tf_file = Path("calibration_image_sample_data_20x128x128x3_float32.npy")
if not onnx2tf_file.exists():
attempt_download_asset(f"{onnx2tf_file}.zip", unzip=True, delete=True)
# Export to TF
images = None
if self.args.int8 and self.args.data:
images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
images = (
torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz)
.permute(0, 2, 3, 1)
.numpy()
.astype(np.float32)
)
# Export to ONNX
if isinstance(self.model.model[-1], RTDETRDecoder):
self.args.opset = self.args.opset or 19
assert 16 <= self.args.opset <= 19, "RTDETR export requires opset>=16;<=19"
self.args.simplify = True
f_onnx = self.export_onnx()
# Export to TF
np_data = None
if self.args.int8:
tmp_file = f / "tmp_tflite_int8_calibration_images.npy" # int8 calibration images file
if self.args.data:
f.mkdir()
images = [batch["img"] for batch in self.get_int8_calibration_dataloader(prefix)]
images = torch.nn.functional.interpolate(torch.cat(images, 0).float(), size=self.imgsz).permute(
0, 2, 3, 1
)
np.save(str(tmp_file), images.numpy().astype(np.float32)) # BHWC
np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
import onnx2tf # scoped for after ONNX export for reduced conflict during import
LOGGER.info(f"{prefix} starting TFLite export with onnx2tf {onnx2tf.__version__}...")
keras_model = onnx2tf.convert(
input_onnx_file_path=f_onnx,
output_folder_path=str(f),
not_use_onnxsim=True,
verbosity="error", # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
output_integer_quantized_tflite=self.args.int8,
custom_input_op_name_np_data_path=np_data,
enable_batchmatmul_unfold=True and not self.args.int8, # fix lower no. of detected objects on GPU delegate
output_signaturedefs=True, # fix error with Attention block group convolution
disable_group_convolution=self.args.format in {"tfjs", "edgetpu"}, # fix error with group convolution
f_onnx = self.export_onnx() # ensure ONNX is available
keras_model = onnx2saved_model(
f_onnx,
f,
int8=self.args.int8,
images=images,
disable_group_convolution=self.args.format in {"tfjs", "edgetpu"},
prefix=prefix,
)
YAML.save(f / "metadata.yaml", self.metadata) # add metadata.yaml
# Remove/rename TFLite models
if self.args.int8:
tmp_file.unlink(missing_ok=True)
for file in f.rglob("*_dynamic_range_quant.tflite"):
file.rename(file.with_name(file.stem.replace("_dynamic_range_quant", "_int8") + file.suffix))
for file in f.rglob("*_integer_quant_with_int16_act.tflite"):
file.unlink() # delete extra fp16 activation TFLite files
# Add TFLite metadata
for file in f.rglob("*.tflite"):
f.unlink() if "quant_with_int16_act.tflite" in str(f) else self._add_tflite_metadata(file)
file.unlink() if "quant_with_int16_act.tflite" in str(file) else self._add_tflite_metadata(file)
return str(f), keras_model # or keras_model = tf.saved_model.load(f, tags=None, options=None)
@try_export
def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
"""Export YOLO model to TensorFlow GraphDef *.pb format https://github.com/leimao/Frozen-Graph-TensorFlow."""
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
f = self.file.with_suffix(".pb")
m = tf.function(lambda x: keras_model(x)) # full model
m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
frozen_func = convert_variables_to_constants_v2(m)
frozen_func.graph.as_graph_def()
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
keras2pb(keras_model, f, prefix)
return f
@try_export
@ -1189,22 +1160,11 @@ class Exporter:
"sudo apt-get install edgetpu-compiler",
):
subprocess.run(c if is_sudo_available() else c.replace("sudo ", ""), shell=True, check=True)
ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]
LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
tflite2edgetpu(tflite_file=tflite_model, output_dir=tflite_model.parent, prefix=prefix)
f = str(tflite_model).replace(".tflite", "_edgetpu.tflite") # Edge TPU model
cmd = (
"edgetpu_compiler "
f'--out_dir "{Path(f).parent}" '
"--show_operations "
"--search_delegate "
"--delegate_search_step 30 "
"--timeout_sec 180 "
f'"{tflite_model}"'
)
LOGGER.info(f"{prefix} running '{cmd}'")
subprocess.run(cmd, shell=True)
self._add_tflite_metadata(f)
return f
@ -1212,31 +1172,10 @@ class Exporter:
def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
"""Export YOLO model to TensorFlow.js format."""
check_requirements("tensorflowjs")
import tensorflow as tf
import tensorflowjs as tfjs
LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
f = str(self.file).replace(self.file.suffix, "_web_model") # js dir
f_pb = str(self.file.with_suffix(".pb")) # *.pb path
gd = tf.Graph().as_graph_def() # TF GraphDef
with open(f_pb, "rb") as file:
gd.ParseFromString(file.read())
outputs = ",".join(gd_outputs(gd))
LOGGER.info(f"\n{prefix} output node names: {outputs}")
quantization = "--quantize_float16" if self.args.half else "--quantize_uint8" if self.args.int8 else ""
with spaces_in_path(f_pb) as fpb_, spaces_in_path(f) as f_: # exporter can not handle spaces in path
cmd = (
"tensorflowjs_converter "
f'--input_format=tf_frozen_model {quantization} --output_node_names={outputs} "{fpb_}" "{f_}"'
)
LOGGER.info(f"{prefix} running '{cmd}'")
subprocess.run(cmd, shell=True)
if " " in f:
LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{f}'.")
pb2tfjs(pb_file=f_pb, output_dir=f, half=self.args.half, int8=self.args.int8, prefix=prefix)
# Add metadata
YAML.save(Path(f) / "metadata.yaml", self.metadata) # add metadata.yaml
return f

2
ultralytics/nn/autobackend.py
View file

@ -428,7 +428,7 @@ class AutoBackend(nn.Module):
LOGGER.info(f"Loading {w} for TensorFlow GraphDef inference...")
import tensorflow as tf
from ultralytics.engine.exporter import gd_outputs
from ultralytics.utils.export.tensorflow import gd_outputs
def wrap_frozen_graph(gd, inputs, outputs):
"""Wrap frozen graphs for deployment."""

35
ultralytics/nn/modules/head.py
View file

@ -166,22 +166,8 @@ class Detect(nn.Module):
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
if self.export and self.format in {"saved_model", "pb", "tflite", "edgetpu", "tfjs"}: # avoid TF FlexSplitV ops
box = x_cat[:, : self.reg_max * 4]
cls = x_cat[:, self.reg_max * 4 :]
else:
box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
if self.export and self.format in {"tflite", "edgetpu"}:
# Precompute normalization factor to increase numerical stability
# See https://github.com/ultralytics/ultralytics/issues/7371
grid_h = shape[2]
grid_w = shape[3]
grid_size = torch.tensor([grid_w, grid_h, grid_w, grid_h], device=box.device).reshape(1, 4, 1)
norm = self.strides / (self.stride[0] * grid_size)
dbox = self.decode_bboxes(self.dfl(box) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
else:
dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides
box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides
return torch.cat((dbox, cls.sigmoid()), 1)
def bias_init(self):
@ -391,20 +377,9 @@ class Pose(Detect):
"""Decode keypoints from predictions."""
ndim = self.kpt_shape[1]
if self.export:
if self.format in {
"tflite",
"edgetpu",
}: # required for TFLite export to avoid 'PLACEHOLDER_FOR_GREATER_OP_CODES' bug
# Precompute normalization factor to increase numerical stability
y = kpts.view(bs, *self.kpt_shape, -1)
grid_h, grid_w = self.shape[2], self.shape[3]
grid_size = torch.tensor([grid_w, grid_h], device=y.device).reshape(1, 2, 1)
norm = self.strides / (self.stride[0] * grid_size)
a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * norm
else:
# NCNN fix
y = kpts.view(bs, *self.kpt_shape, -1)
a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * self.strides
# NCNN fix
y = kpts.view(bs, *self.kpt_shape, -1)
a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * self.strides
if ndim == 3:
a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
return a.view(bs, self.nk, -1)

243
ultralytics/utils/export/__init__.py
View file

@ -1,242 +1,7 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
from .engine import onnx2engine, torch2onnx
from .imx import torch2imx
from .tensorflow import keras2pb, onnx2saved_model, pb2tfjs, tflite2edgetpu
import json
from pathlib import Path
import torch
from ultralytics.utils import IS_JETSON, LOGGER
from ultralytics.utils.torch_utils import TORCH_2_4
from .imx import torch2imx # noqa
def torch2onnx(
torch_model: torch.nn.Module,
im: torch.Tensor,
onnx_file: str,
opset: int = 14,
input_names: list[str] = ["images"],
output_names: list[str] = ["output0"],
dynamic: bool | dict = False,
) -> None:
"""
Export a PyTorch model to ONNX format.
Args:
torch_model (torch.nn.Module): The PyTorch model to export.
im (torch.Tensor): Example input tensor for the model.
onnx_file (str): Path to save the exported ONNX file.
opset (int): ONNX opset version to use for export.
input_names (list[str]): List of input tensor names.
output_names (list[str]): List of output tensor names.
dynamic (bool | dict, optional): Whether to enable dynamic axes.
Notes:
Setting `do_constant_folding=True` may cause issues with DNN inference for torch>=1.12.
"""
kwargs = {"dynamo": False} if TORCH_2_4 else {}
torch.onnx.export(
torch_model,
im,
onnx_file,
verbose=False,
opset_version=opset,
do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic or None,
**kwargs,
)
def onnx2engine(
onnx_file: str,
engine_file: str | None = None,
workspace: int | None = None,
half: bool = False,
int8: bool = False,
dynamic: bool = False,
shape: tuple[int, int, int, int] = (1, 3, 640, 640),
dla: int | None = None,
dataset=None,
metadata: dict | None = None,
verbose: bool = False,
prefix: str = "",
) -> None:
"""
Export a YOLO model to TensorRT engine format.
Args:
onnx_file (str): Path to the ONNX file to be converted.
engine_file (str, optional): Path to save the generated TensorRT engine file.
workspace (int, optional): Workspace size in GB for TensorRT.
half (bool, optional): Enable FP16 precision.
int8 (bool, optional): Enable INT8 precision.
dynamic (bool, optional): Enable dynamic input shapes.
shape (tuple[int, int, int, int], optional): Input shape (batch, channels, height, width).
dla (int, optional): DLA core to use (Jetson devices only).
dataset (ultralytics.data.build.InfiniteDataLoader, optional): Dataset for INT8 calibration.
metadata (dict, optional): Metadata to include in the engine file.
verbose (bool, optional): Enable verbose logging.
prefix (str, optional): Prefix for log messages.
Raises:
ValueError: If DLA is enabled on non-Jetson devices or required precision is not set.
RuntimeError: If the ONNX file cannot be parsed.
Notes:
TensorRT version compatibility is handled for workspace size and engine building.
INT8 calibration requires a dataset and generates a calibration cache.
Metadata is serialized and written to the engine file if provided.
"""
import tensorrt as trt
engine_file = engine_file or Path(onnx_file).with_suffix(".engine")
logger = trt.Logger(trt.Logger.INFO)
if verbose:
logger.min_severity = trt.Logger.Severity.VERBOSE
# Engine builder
builder = trt.Builder(logger)
config = builder.create_builder_config()
workspace_bytes = int((workspace or 0) * (1 << 30))
is_trt10 = int(trt.__version__.split(".", 1)[0]) >= 10 # is TensorRT >= 10
if is_trt10 and workspace_bytes > 0:
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
elif workspace_bytes > 0: # TensorRT versions 7, 8
config.max_workspace_size = workspace_bytes
flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(flag)
half = builder.platform_has_fast_fp16 and half
int8 = builder.platform_has_fast_int8 and int8
# Optionally switch to DLA if enabled
if dla is not None:
if not IS_JETSON:
raise ValueError("DLA is only available on NVIDIA Jetson devices")
LOGGER.info(f"{prefix} enabling DLA on core {dla}...")
if not half and not int8:
raise ValueError(
"DLA requires either 'half=True' (FP16) or 'int8=True' (INT8) to be enabled. Please enable one of them and try again."
)
config.default_device_type = trt.DeviceType.DLA
config.DLA_core = int(dla)
config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
# Read ONNX file
parser = trt.OnnxParser(network, logger)
if not parser.parse_from_file(onnx_file):
raise RuntimeError(f"failed to load ONNX file: {onnx_file}")
# Network inputs
inputs = [network.get_input(i) for i in range(network.num_inputs)]
outputs = [network.get_output(i) for i in range(network.num_outputs)]
for inp in inputs:
LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
for out in outputs:
LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
if dynamic:
profile = builder.create_optimization_profile()
min_shape = (1, shape[1], 32, 32) # minimum input shape
max_shape = (*shape[:2], *(int(max(2, workspace or 2) * d) for d in shape[2:])) # max input shape
for inp in inputs:
profile.set_shape(inp.name, min=min_shape, opt=shape, max=max_shape)
config.add_optimization_profile(profile)
if int8:
config.set_calibration_profile(profile)
LOGGER.info(f"{prefix} building {'INT8' if int8 else 'FP' + ('16' if half else '32')} engine as {engine_file}")
if int8:
config.set_flag(trt.BuilderFlag.INT8)
config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
class EngineCalibrator(trt.IInt8Calibrator):
"""
Custom INT8 calibrator for TensorRT engine optimization.
This calibrator provides the necessary interface for TensorRT to perform INT8 quantization calibration
using a dataset. It handles batch generation, caching, and calibration algorithm selection.
Attributes:
dataset: Dataset for calibration.
data_iter: Iterator over the calibration dataset.
algo (trt.CalibrationAlgoType): Calibration algorithm type.
batch (int): Batch size for calibration.
cache (Path): Path to save the calibration cache.
Methods:
get_algorithm: Get the calibration algorithm to use.
get_batch_size: Get the batch size to use for calibration.
get_batch: Get the next batch to use for calibration.
read_calibration_cache: Use existing cache instead of calibrating again.
write_calibration_cache: Write calibration cache to disk.
"""
def __init__(
self,
dataset, # ultralytics.data.build.InfiniteDataLoader
cache: str = "",
) -> None:
"""Initialize the INT8 calibrator with dataset and cache path."""
trt.IInt8Calibrator.__init__(self)
self.dataset = dataset
self.data_iter = iter(dataset)
self.algo = (
trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2 # DLA quantization needs ENTROPY_CALIBRATION_2
if dla is not None
else trt.CalibrationAlgoType.MINMAX_CALIBRATION
)
self.batch = dataset.batch_size
self.cache = Path(cache)
def get_algorithm(self) -> trt.CalibrationAlgoType:
"""Get the calibration algorithm to use."""
return self.algo
def get_batch_size(self) -> int:
"""Get the batch size to use for calibration."""
return self.batch or 1
def get_batch(self, names) -> list[int] | None:
"""Get the next batch to use for calibration, as a list of device memory pointers."""
try:
im0s = next(self.data_iter)["img"] / 255.0
im0s = im0s.to("cuda") if im0s.device.type == "cpu" else im0s
return [int(im0s.data_ptr())]
except StopIteration:
# Return None to signal to TensorRT there is no calibration data remaining
return None
def read_calibration_cache(self) -> bytes | None:
"""Use existing cache instead of calibrating again, otherwise, implicitly return None."""
if self.cache.exists() and self.cache.suffix == ".cache":
return self.cache.read_bytes()
def write_calibration_cache(self, cache: bytes) -> None:
"""Write calibration cache to disk."""
_ = self.cache.write_bytes(cache)
# Load dataset w/ builder (for batching) and calibrate
config.int8_calibrator = EngineCalibrator(
dataset=dataset,
cache=str(Path(onnx_file).with_suffix(".cache")),
)
elif half:
config.set_flag(trt.BuilderFlag.FP16)
# Write file
build = builder.build_serialized_network if is_trt10 else builder.build_engine
with build(network, config) as engine, open(engine_file, "wb") as t:
# Metadata
if metadata is not None:
meta = json.dumps(metadata)
t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
t.write(meta.encode())
# Model
t.write(engine if is_trt10 else engine.serialize())
__all__ = ["keras2pb", "onnx2engine", "onnx2saved_model", "pb2tfjs", "tflite2edgetpu", "torch2imx", "torch2onnx"]

240
ultralytics/utils/export/engine.py Normal file
View file

@ -0,0 +1,240 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
import json
from pathlib import Path
import torch
from ultralytics.utils import IS_JETSON, LOGGER
from ultralytics.utils.torch_utils import TORCH_2_4
def torch2onnx(
torch_model: torch.nn.Module,
im: torch.Tensor,
onnx_file: str,
opset: int = 14,
input_names: list[str] = ["images"],
output_names: list[str] = ["output0"],
dynamic: bool | dict = False,
) -> None:
"""
Export a PyTorch model to ONNX format.
Args:
torch_model (torch.nn.Module): The PyTorch model to export.
im (torch.Tensor): Example input tensor for the model.
onnx_file (str): Path to save the exported ONNX file.
opset (int): ONNX opset version to use for export.
input_names (list[str]): List of input tensor names.
output_names (list[str]): List of output tensor names.
dynamic (bool | dict, optional): Whether to enable dynamic axes.
Notes:
Setting `do_constant_folding=True` may cause issues with DNN inference for torch>=1.12.
"""
kwargs = {"dynamo": False} if TORCH_2_4 else {}
torch.onnx.export(
torch_model,
im,
onnx_file,
verbose=False,
opset_version=opset,
do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic or None,
**kwargs,
)
def onnx2engine(
onnx_file: str,
engine_file: str | None = None,
workspace: int | None = None,
half: bool = False,
int8: bool = False,
dynamic: bool = False,
shape: tuple[int, int, int, int] = (1, 3, 640, 640),
dla: int | None = None,
dataset=None,
metadata: dict | None = None,
verbose: bool = False,
prefix: str = "",
) -> None:
"""
Export a YOLO model to TensorRT engine format.
Args:
onnx_file (str): Path to the ONNX file to be converted.
engine_file (str, optional): Path to save the generated TensorRT engine file.
workspace (int, optional): Workspace size in GB for TensorRT.
half (bool, optional): Enable FP16 precision.
int8 (bool, optional): Enable INT8 precision.
dynamic (bool, optional): Enable dynamic input shapes.
shape (tuple[int, int, int, int], optional): Input shape (batch, channels, height, width).
dla (int, optional): DLA core to use (Jetson devices only).
dataset (ultralytics.data.build.InfiniteDataLoader, optional): Dataset for INT8 calibration.
metadata (dict, optional): Metadata to include in the engine file.
verbose (bool, optional): Enable verbose logging.
prefix (str, optional): Prefix for log messages.
Raises:
ValueError: If DLA is enabled on non-Jetson devices or required precision is not set.
RuntimeError: If the ONNX file cannot be parsed.
Notes:
TensorRT version compatibility is handled for workspace size and engine building.
INT8 calibration requires a dataset and generates a calibration cache.
Metadata is serialized and written to the engine file if provided.
"""
import tensorrt as trt
engine_file = engine_file or Path(onnx_file).with_suffix(".engine")
logger = trt.Logger(trt.Logger.INFO)
if verbose:
logger.min_severity = trt.Logger.Severity.VERBOSE
# Engine builder
builder = trt.Builder(logger)
config = builder.create_builder_config()
workspace_bytes = int((workspace or 0) * (1 << 30))
is_trt10 = int(trt.__version__.split(".", 1)[0]) >= 10 # is TensorRT >= 10
if is_trt10 and workspace_bytes > 0:
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
elif workspace_bytes > 0: # TensorRT versions 7, 8
config.max_workspace_size = workspace_bytes
flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(flag)
half = builder.platform_has_fast_fp16 and half
int8 = builder.platform_has_fast_int8 and int8
# Optionally switch to DLA if enabled
if dla is not None:
if not IS_JETSON:
raise ValueError("DLA is only available on NVIDIA Jetson devices")
LOGGER.info(f"{prefix} enabling DLA on core {dla}...")
if not half and not int8:
raise ValueError(
"DLA requires either 'half=True' (FP16) or 'int8=True' (INT8) to be enabled. Please enable one of them and try again."
)
config.default_device_type = trt.DeviceType.DLA
config.DLA_core = int(dla)
config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
# Read ONNX file
parser = trt.OnnxParser(network, logger)
if not parser.parse_from_file(onnx_file):
raise RuntimeError(f"failed to load ONNX file: {onnx_file}")
# Network inputs
inputs = [network.get_input(i) for i in range(network.num_inputs)]
outputs = [network.get_output(i) for i in range(network.num_outputs)]
for inp in inputs:
LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
for out in outputs:
LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
if dynamic:
profile = builder.create_optimization_profile()
min_shape = (1, shape[1], 32, 32) # minimum input shape
max_shape = (*shape[:2], *(int(max(2, workspace or 2) * d) for d in shape[2:])) # max input shape
for inp in inputs:
profile.set_shape(inp.name, min=min_shape, opt=shape, max=max_shape)
config.add_optimization_profile(profile)
if int8:
config.set_calibration_profile(profile)
LOGGER.info(f"{prefix} building {'INT8' if int8 else 'FP' + ('16' if half else '32')} engine as {engine_file}")
if int8:
config.set_flag(trt.BuilderFlag.INT8)
config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
class EngineCalibrator(trt.IInt8Calibrator):
"""
Custom INT8 calibrator for TensorRT engine optimization.
This calibrator provides the necessary interface for TensorRT to perform INT8 quantization calibration
using a dataset. It handles batch generation, caching, and calibration algorithm selection.
Attributes:
dataset: Dataset for calibration.
data_iter: Iterator over the calibration dataset.
algo (trt.CalibrationAlgoType): Calibration algorithm type.
batch (int): Batch size for calibration.
cache (Path): Path to save the calibration cache.
Methods:
get_algorithm: Get the calibration algorithm to use.
get_batch_size: Get the batch size to use for calibration.
get_batch: Get the next batch to use for calibration.
read_calibration_cache: Use existing cache instead of calibrating again.
write_calibration_cache: Write calibration cache to disk.
"""
def __init__(
self,
dataset, # ultralytics.data.build.InfiniteDataLoader
cache: str = "",
) -> None:
"""Initialize the INT8 calibrator with dataset and cache path."""
trt.IInt8Calibrator.__init__(self)
self.dataset = dataset
self.data_iter = iter(dataset)
self.algo = (
trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2 # DLA quantization needs ENTROPY_CALIBRATION_2
if dla is not None
else trt.CalibrationAlgoType.MINMAX_CALIBRATION
)
self.batch = dataset.batch_size
self.cache = Path(cache)
def get_algorithm(self) -> trt.CalibrationAlgoType:
"""Get the calibration algorithm to use."""
return self.algo
def get_batch_size(self) -> int:
"""Get the batch size to use for calibration."""
return self.batch or 1
def get_batch(self, names) -> list[int] | None:
"""Get the next batch to use for calibration, as a list of device memory pointers."""
try:
im0s = next(self.data_iter)["img"] / 255.0
im0s = im0s.to("cuda") if im0s.device.type == "cpu" else im0s
return [int(im0s.data_ptr())]
except StopIteration:
# Return None to signal to TensorRT there is no calibration data remaining
return None
def read_calibration_cache(self) -> bytes | None:
"""Use existing cache instead of calibrating again, otherwise, implicitly return None."""
if self.cache.exists() and self.cache.suffix == ".cache":
return self.cache.read_bytes()
def write_calibration_cache(self, cache: bytes) -> None:
"""Write calibration cache to disk."""
_ = self.cache.write_bytes(cache)
# Load dataset w/ builder (for batching) and calibrate
config.int8_calibrator = EngineCalibrator(
dataset=dataset,
cache=str(Path(onnx_file).with_suffix(".cache")),
)
elif half:
config.set_flag(trt.BuilderFlag.FP16)
# Write file
build = builder.build_serialized_network if is_trt10 else builder.build_engine
with build(network, config) as engine, open(engine_file, "wb") as t:
# Metadata
if metadata is not None:
meta = json.dumps(metadata)
t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
t.write(meta.encode())
# Model
t.write(engine if is_trt10 else engine.serialize())

221
ultralytics/utils/export/tensorflow.py Normal file
View file

@ -0,0 +1,221 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from __future__ import annotations
from pathlib import Path
import numpy as np
import torch
from ultralytics.nn.modules import Detect, Pose
from ultralytics.utils import LOGGER
from ultralytics.utils.downloads import attempt_download_asset
from ultralytics.utils.files import spaces_in_path
from ultralytics.utils.tal import make_anchors
def tf_wrapper(model: torch.nn.Module) -> torch.nn.Module:
"""A wrapper to add TensorFlow compatible inference methods to Detect and Pose layers."""
for m in model.modules():
if not isinstance(m, Detect):
continue
import types
m._inference = types.MethodType(_tf_inference, m)
if type(m) is Pose:
m.kpts_decode = types.MethodType(tf_kpts_decode, m)
return model
def _tf_inference(self, x: list[torch.Tensor]) -> tuple[torch.Tensor]:
"""Decode boxes and cls scores for tf object detection."""
shape = x[0].shape # BCHW
x_cat = torch.cat([xi.view(x[0].shape[0], self.no, -1) for xi in x], 2)
box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
if self.dynamic or self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
grid_h, grid_w = shape[2], shape[3]
grid_size = torch.tensor([grid_w, grid_h, grid_w, grid_h], device=box.device).reshape(1, 4, 1)
norm = self.strides / (self.stride[0] * grid_size)
dbox = self.decode_bboxes(self.dfl(box) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
return torch.cat((dbox, cls.sigmoid()), 1)
def tf_kpts_decode(self, bs: int, kpts: torch.Tensor) -> torch.Tensor:
"""Decode keypoints for tf pose estimation."""
ndim = self.kpt_shape[1]
# required for TFLite export to avoid 'PLACEHOLDER_FOR_GREATER_OP_CODES' bug
# Precompute normalization factor to increase numerical stability
y = kpts.view(bs, *self.kpt_shape, -1)
grid_h, grid_w = self.shape[2], self.shape[3]
grid_size = torch.tensor([grid_w, grid_h], device=y.device).reshape(1, 2, 1)
norm = self.strides / (self.stride[0] * grid_size)
a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * norm
if ndim == 3:
a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
return a.view(bs, self.nk, -1)
def onnx2saved_model(
onnx_file: str,
output_dir: Path,
int8: bool = False,
images: np.ndarray = None,
disable_group_convolution: bool = False,
prefix="",
):
"""
Convert a ONNX model to TensorFlow SavedModel format via ONNX.
Args:
onnx_file (str): ONNX file path.
output_dir (Path): Output directory path for the SavedModel.
int8 (bool, optional): Enable INT8 quantization. Defaults to False.
images (np.ndarray, optional): Calibration images for INT8 quantization in BHWC format.
disable_group_convolution (bool, optional): Disable group convolution optimization. Defaults to False.
prefix (str, optional): Logging prefix. Defaults to "".
Returns:
(keras.Model): Converted Keras model.
Note:
Requires onnx2tf package. Downloads calibration data if INT8 quantization is enabled.
Removes temporary files and renames quantized models after conversion.
"""
# Pre-download calibration file to fix https://github.com/PINTO0309/onnx2tf/issues/545
onnx2tf_file = Path("calibration_image_sample_data_20x128x128x3_float32.npy")
if not onnx2tf_file.exists():
attempt_download_asset(f"{onnx2tf_file}.zip", unzip=True, delete=True)
np_data = None
if int8:
tmp_file = output_dir / "tmp_tflite_int8_calibration_images.npy" # int8 calibration images file
if images is not None:
output_dir.mkdir()
np.save(str(tmp_file), images) # BHWC
np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
import onnx2tf # scoped for after ONNX export for reduced conflict during import
LOGGER.info(f"{prefix} starting TFLite export with onnx2tf {onnx2tf.__version__}...")
keras_model = onnx2tf.convert(
input_onnx_file_path=onnx_file,
output_folder_path=str(output_dir),
not_use_onnxsim=True,
verbosity="error", # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
output_integer_quantized_tflite=int8,
custom_input_op_name_np_data_path=np_data,
enable_batchmatmul_unfold=True and not int8, # fix lower no. of detected objects on GPU delegate
output_signaturedefs=True, # fix error with Attention block group convolution
disable_group_convolution=disable_group_convolution, # fix error with group convolution
)
# Remove/rename TFLite models
if int8:
tmp_file.unlink(missing_ok=True)
for file in output_dir.rglob("*_dynamic_range_quant.tflite"):
file.rename(file.with_name(file.stem.replace("_dynamic_range_quant", "_int8") + file.suffix))
for file in output_dir.rglob("*_integer_quant_with_int16_act.tflite"):
file.unlink() # delete extra fp16 activation TFLite files
return keras_model
def keras2pb(keras_model, file: Path, prefix=""):
"""
Convert a Keras model to TensorFlow GraphDef (.pb) format.
Args:
keras_model(tf_keras): Keras model to convert to frozen graph format.
file (Path): Output file path (suffix will be changed to .pb).
prefix (str, optional): Logging prefix. Defaults to "".
Note:
Creates a frozen graph by converting variables to constants for inference optimization.
"""
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
m = tf.function(lambda x: keras_model(x)) # full model
m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
frozen_func = convert_variables_to_constants_v2(m)
frozen_func.graph.as_graph_def()
tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(file.parent), name=file.name, as_text=False)
def tflite2edgetpu(tflite_file: str | Path, output_dir: str | Path, prefix: str = ""):
"""
Convert a TensorFlow Lite model to Edge TPU format using the Edge TPU compiler.
Args:
tflite_file (str | Path): Path to the input TensorFlow Lite (.tflite) model file.
output_dir (str | Path): Output directory path for the compiled Edge TPU model.
prefix (str, optional): Logging prefix. Defaults to "".
Note:
Requires the Edge TPU compiler to be installed. The function compiles the TFLite model
for optimal performance on Google's Edge TPU hardware accelerator.
"""
import subprocess
cmd = (
"edgetpu_compiler "
f'--out_dir "{output_dir}" '
"--show_operations "
"--search_delegate "
"--delegate_search_step 30 "
"--timeout_sec 180 "
f'"{tflite_file}"'
)
LOGGER.info(f"{prefix} running '{cmd}'")
subprocess.run(cmd, shell=True)
def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = False, prefix: str = ""):
"""
Convert a TensorFlow GraphDef (.pb) model to TensorFlow.js format.
Args:
pb_file (str): Path to the input TensorFlow GraphDef (.pb) model file.
output_dir (str): Output directory path for the converted TensorFlow.js model.
half (bool, optional): Enable FP16 quantization. Defaults to False.
int8 (bool, optional): Enable INT8 quantization. Defaults to False.
prefix (str, optional): Logging prefix. Defaults to "".
Note:
Requires tensorflowjs package. Uses tensorflowjs_converter command-line tool for conversion.
Handles spaces in file paths and warns if output directory contains spaces.
"""
import subprocess
import tensorflow as tf
import tensorflowjs as tfjs
LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
gd = tf.Graph().as_graph_def() # TF GraphDef
with open(pb_file, "rb") as file:
gd.ParseFromString(file.read())
outputs = ",".join(gd_outputs(gd))
LOGGER.info(f"\n{prefix} output node names: {outputs}")
quantization = "--quantize_float16" if half else "--quantize_uint8" if int8 else ""
with spaces_in_path(pb_file) as fpb_, spaces_in_path(output_dir) as f_: # exporter can not handle spaces in path
cmd = (
"tensorflowjs_converter "
f'--input_format=tf_frozen_model {quantization} --output_node_names={outputs} "{fpb_}" "{f_}"'
)
LOGGER.info(f"{prefix} running '{cmd}'")
subprocess.run(cmd, shell=True)
if " " in output_dir:
LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{output_dir}'.")
def gd_outputs(gd):
"""Return TensorFlow GraphDef model output node names."""
name_list, input_list = [], []
for node in gd.node: # tensorflow.core.framework.node_def_pb2.NodeDef
name_list.append(node.name)
input_list.extend(node.input)
return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))