单目3D初始代码

ultralytics/utils/export/executorch.py

@@ -0,0 +1,79 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from functools import partial
+from pathlib import Path
+
+import torch
+
+from ultralytics.nn.modules import Pose, Pose26
+from ultralytics.utils import LOGGER, YAML
+
+
+def executorch_wrapper(model: torch.nn.Module) -> torch.nn.Module:
+    """Apply ExecuTorch-specific model patches required for export/runtime compatibility."""
+    import types
+
+    for m in model.modules():
+        if not isinstance(m, Pose):
+            continue
+        m.kpts_decode = types.MethodType(partial(_executorch_kpts_decode, is_pose26=type(m) is Pose26), m)
+    return model
+
+
+def _executorch_kpts_decode(self, kpts: torch.Tensor, is_pose26: bool = False) -> torch.Tensor:
+    """Decode pose keypoints for ExecuTorch export with XNNPACK-safe broadcasting."""
+    ndim = self.kpt_shape[1]
+    bs = kpts.shape[0]
+    y = kpts.view(bs, *self.kpt_shape, -1)
+
+    # XNNPACK requires explicit dim matching for broadcasting, expand 2D tensors to 4D.
+    anchors = self.anchors[None, None]
+    strides = self.strides[None, None]
+    a = ((y[:, :, :2] + anchors) if is_pose26 else (y[:, :, :2] * 2.0 + (anchors - 0.5))) * strides
+    if ndim == 3:
+        a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
+    return a.view(bs, self.nk, -1)
+
+
+def torch2executorch(
+    model: torch.nn.Module,
+    file: Path | str,
+    sample_input: torch.Tensor,
+    metadata: dict | None = None,
+    prefix: str = "",
+) -> str:
+    """Export a PyTorch model to ExecuTorch format.
+
+    Args:
+        model (torch.nn.Module): The PyTorch model to export.
+        file (Path | str): Source model file path used to derive output names.
+        sample_input (torch.Tensor): Example input tensor for tracing/export.
+        metadata (dict | None, optional): Optional metadata to save as YAML.
+        prefix (str, optional): Prefix for log messages.
+
+    Returns:
+        (str): Path to the exported ExecuTorch model directory.
+    """
+    from executorch import version as executorch_version
+    from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+    from executorch.exir import to_edge_transform_and_lower
+
+    LOGGER.info(f"\n{prefix} starting export with ExecuTorch {executorch_version.__version__}...")
+
+    file = Path(file)
+    output_dir = Path(str(file).replace(file.suffix, "_executorch_model"))
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    pte_file = output_dir / file.with_suffix(".pte").name
+    et_program = to_edge_transform_and_lower(
+        torch.export.export(model, (sample_input,)),
+        partitioner=[XnnpackPartitioner()],
+    ).to_executorch()
+    pte_file.write_bytes(et_program.buffer)
+
+    if metadata is not None:
+        YAML.save(output_dir / "metadata.yaml", metadata)
+
+    return str(output_dir)