dgenerate-ultralytics-headless 8.3.221__py3-none-any.whl → 8.3.223__py3-none-any.whl

ultralytics/nn/modules/head.py

@@ -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)

ultralytics/utils/__init__.py

@@ -795,13 +795,13 @@ def is_pip_package(filepath: str = __name__) -> bool:
 
 def is_dir_writeable(dir_path: str | Path) -> bool:
     """
-    Check if a directory is writeable.
+    Check if a directory is writable.
 
     Args:
         dir_path (str | Path): The path to the directory.
 
     Returns:
-        (bool): True if the directory is writeable, False otherwise.
+        (bool): True if the directory is writable, False otherwise.
     """
     return os.access(str(dir_path), os.W_OK)
 
@@ -882,14 +882,14 @@ def get_user_config_dir(sub_dir="Ultralytics"):
         p.mkdir(parents=True, exist_ok=True)
         return p
 
-    # Fallbacks for Docker, GCP/AWS functions where only /tmp is writeable
+    # Fallbacks for Docker, GCP/AWS functions where only /tmp is writable
     for alt in [Path("/tmp") / sub_dir, Path.cwd() / sub_dir]:
         if alt.exists():
             return alt
         if is_dir_writeable(alt.parent):
             alt.mkdir(parents=True, exist_ok=True)
             LOGGER.warning(
-                f"user config directory '{p}' is not writeable, using '{alt}'. Set YOLO_CONFIG_DIR to override."
+                f"user config directory '{p}' is not writable, using '{alt}'. Set YOLO_CONFIG_DIR to override."
             )
             return alt
 

ultralytics/utils/benchmarks.py

@@ -144,7 +144,9 @@ def benchmark(
             if format == "imx":
                 assert not is_end2end
                 assert not isinstance(model, YOLOWorld), "YOLOWorldv2 IMX exports not supported"
-                assert model.task == "detect", "IMX only supported for detection task"
+                assert model.task in {"detect", "classify", "pose"}, (
+                    "IMX export is only supported for detection, classification and pose estimation tasks"
+                )
                 assert "C2f" in model.__str__(), "IMX only supported for YOLOv8n and YOLO11n"
             if format == "rknn":
                 assert not isinstance(model, YOLOWorld), "YOLOWorldv2 RKNN exports not supported yet"

ultralytics/utils/export/__init__.py

@@ -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"]

ultralytics/utils/export/engine.py

@@ -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())