Reference for ultralytics/nn/backends/axelera.py#

class AxeleraBackend(BaseBackend):
    """Axelera AI inference backend for Axelera Metis AI accelerators.

    Loads compiled Axelera models (.axm files) and runs inference using the Axelera AI runtime SDK.
    """

def forward(self, im: torch.Tensor) -> list:
    """Run inference on the Axelera hardware accelerator.

    Args:
        im (torch.Tensor): Input image tensor in BCHW format, normalized to [0, 1].

    Returns:
        (list): Model predictions as a list of output arrays.
    """
    return self.model(im.cpu())

def load_model(self, weight: str | Path) -> None:
    """Load an Axelera model from a directory containing a .axm file.

    Args:
        weight (str | Path): Path to the Axelera model directory containing the .axm binary.
    """
    try:
        from axelera.runtime import op
    except ImportError:
        check_requirements(
            "axelera-rt==1.7.0",
            cmds="--extra-index-url https://software.axelera.ai/artifactory/api/pypi/axelera-pypi/simple",
        )

    from axelera.runtime import op

    w = Path(weight)
    found = next(w.rglob("*.axm"), None)
    if found is None:
        raise FileNotFoundError(f"No .axm file found in: {w}")

    self.model = op.load(str(found)).optimized()

    # Load metadata
    metadata_file = found.parent / "metadata.yaml"
    if metadata_file.exists():
        from ultralytics.utils import YAML

        self.apply_metadata(YAML.load(metadata_file))