Reference for ultralytics/utils/export/imx.py

def _inference(self, x: list[torch.Tensor]) -> tuple[torch.Tensor]:
    """Decode boxes and cls scores for imx object detection."""
    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)
    dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides
    return dbox.transpose(1, 2), cls.sigmoid().permute(0, 2, 1)

def pose_forward(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward pass for imx pose estimation, including keypoint decoding."""
    bs = x[0].shape[0]  # batch size
    kpt = torch.cat([self.cv4[i](x[i]).view(bs, self.nk, -1) for i in range(self.nl)], -1)  # (bs, 17*3, h*w)
    x = Detect.forward(self, x)
    pred_kpt = self.kpts_decode(bs, kpt)
    return (*x, pred_kpt.permute(0, 2, 1))

def torch2imx(
    model: torch.nn.Module,
    file: Path | str,
    conf: float,
    iou: float,
    max_det: int,
    metadata: dict | None = None,
    gptq: bool = False,
    dataset=None,
    prefix: str = "",
):
    """
    Export YOLO model to IMX format for deployment on Sony IMX500 devices.

    This function quantizes a YOLO model using Model Compression Toolkit (MCT) and exports it
    to IMX format compatible with Sony IMX500 edge devices. It supports both YOLOv8n and YOLO11n
    models for detection and pose estimation tasks.

    Args:
        model (torch.nn.Module): The YOLO model to export. Must be YOLOv8n or YOLO11n.
        file (Path | str): Output file path for the exported model.
        conf (float): Confidence threshold for NMS post-processing.
        iou (float): IoU threshold for NMS post-processing.
        max_det (int): Maximum number of detections to return.
        metadata (dict | None, optional): Metadata to embed in the ONNX model. Defaults to None.
        gptq (bool, optional): Whether to use Gradient-Based Post Training Quantization.
            If False, uses standard Post Training Quantization. Defaults to False.
        dataset (optional): Representative dataset for quantization calibration. Defaults to None.
        prefix (str, optional): Logging prefix string. Defaults to "".

    Returns:
        f (Path): Path to the exported IMX model directory

    Raises:
        ValueError: If the model is not a supported YOLOv8n or YOLO11n variant.

    Example:
        >>> from ultralytics import YOLO
        >>> model = YOLO("yolo11n.pt")
        >>> path, _ = export_imx(model, "model.imx", conf=0.25, iou=0.45, max_det=300)

    Note:
        - Requires model_compression_toolkit, onnx, edgemdt_tpc, and sony_custom_layers packages
        - Only supports YOLOv8n and YOLO11n models (detection and pose tasks)
        - Output includes quantized ONNX model, IMX binary, and labels.txt file
    """
    import model_compression_toolkit as mct
    import onnx
    from edgemdt_tpc import get_target_platform_capabilities

    LOGGER.info(f"\n{prefix} starting export with model_compression_toolkit {mct.__version__}...")

    def representative_dataset_gen(dataloader=dataset):
        for batch in dataloader:
            img = batch["img"]
            img = img / 255.0
            yield [img]

    tpc = get_target_platform_capabilities(tpc_version="4.0", device_type="imx500")

    bit_cfg = mct.core.BitWidthConfig()
    if "C2PSA" in model.__str__():  # YOLO11
        if model.task == "detect":
            layer_names = ["sub", "mul_2", "add_14", "cat_21"]
            weights_memory = 2585350.2439
            n_layers = 238  # 238 layers for fused YOLO11n
        elif model.task == "pose":
            layer_names = ["sub", "mul_2", "add_14", "cat_22", "cat_23", "mul_4", "add_15"]
            weights_memory = 2437771.67
            n_layers = 257  # 257 layers for fused YOLO11n-pose
    else:  # YOLOv8
        if model.task == "detect":
            layer_names = ["sub", "mul", "add_6", "cat_17"]
            weights_memory = 2550540.8
            n_layers = 168  # 168 layers for fused YOLOv8n
        elif model.task == "pose":
            layer_names = ["add_7", "mul_2", "cat_19", "mul", "sub", "add_6", "cat_18"]
            weights_memory = 2482451.85
            n_layers = 187  # 187 layers for fused YOLO11n-pose

    # Check if the model has the expected number of layers
    if len(list(model.modules())) != n_layers:
        raise ValueError("IMX export only supported for YOLOv8n and YOLO11n models.")

    for layer_name in layer_names:
        bit_cfg.set_manual_activation_bit_width([mct.core.common.network_editors.NodeNameFilter(layer_name)], 16)

    config = mct.core.CoreConfig(
        mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=10),
        quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True),
        bit_width_config=bit_cfg,
    )

    resource_utilization = mct.core.ResourceUtilization(weights_memory=weights_memory)

    quant_model = (
        mct.gptq.pytorch_gradient_post_training_quantization(  # Perform Gradient-Based Post Training Quantization
            model=model,
            representative_data_gen=representative_dataset_gen,
            target_resource_utilization=resource_utilization,
            gptq_config=mct.gptq.get_pytorch_gptq_config(
                n_epochs=1000, use_hessian_based_weights=False, use_hessian_sample_attention=False
            ),
            core_config=config,
            target_platform_capabilities=tpc,
        )[0]
        if gptq
        else mct.ptq.pytorch_post_training_quantization(  # Perform post training quantization
            in_module=model,
            representative_data_gen=representative_dataset_gen,
            target_resource_utilization=resource_utilization,
            core_config=config,
            target_platform_capabilities=tpc,
        )[0]
    )

    quant_model = NMSWrapper(
        model=quant_model,
        score_threshold=conf or 0.001,
        iou_threshold=iou,
        max_detections=max_det,
        task=model.task,
    )

    f = Path(str(file).replace(file.suffix, "_imx_model"))
    f.mkdir(exist_ok=True)
    onnx_model = f / Path(str(file.name).replace(file.suffix, "_imx.onnx"))  # js dir
    mct.exporter.pytorch_export_model(
        model=quant_model, save_model_path=onnx_model, repr_dataset=representative_dataset_gen
    )

    model_onnx = onnx.load(onnx_model)  # load onnx model
    for k, v in metadata.items():
        meta = model_onnx.metadata_props.add()
        meta.key, meta.value = k, str(v)

    onnx.save(model_onnx, onnx_model)

    subprocess.run(
        ["imxconv-pt", "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
        check=True,
    )

    # Needed for imx models.
    with open(f / "labels.txt", "w", encoding="utf-8") as file:
        file.writelines([f"{name}\n" for _, name in model.names.items()])

    return f