Horatiu George TodoranandCodex
f418967d55
* perf(ml): share RKNN weights across NPU cores and halve model RAM * perf(ml): default rknn threads to npu cores and cache shape setup * perf(ml): serve rknn runs in arrival order (cherry picked from commit e659fb62f2bb3f6b1db08ed643d3d054a539c415)
80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import orjson
|
|
from numpy.typing import NDArray
|
|
|
|
from immich_ml.config import log, settings
|
|
from immich_ml.models.constants import RKNN_NPU_CORES
|
|
from immich_ml.schemas import ModelInput, ModelTensor, SessionNode, Shape
|
|
|
|
from .rknnpool import RknnNode, RknnPoolExecutor, is_available, soc_name
|
|
|
|
is_available = is_available and settings.rknn
|
|
model_prefix = Path("rknpu") / soc_name if is_available and soc_name is not None else None
|
|
|
|
|
|
def model_path(model_dir: Path, variant: str = "") -> Path:
|
|
return (model_dir / model_prefix if model_prefix else model_dir) / variant / "model.rknn"
|
|
|
|
|
|
def input_layout(compiled: tuple[int, ...], array: ModelTensor) -> str | None:
|
|
if array.ndim != 4:
|
|
return None # token ids carry no layout
|
|
if array.shape[3] == compiled[3]:
|
|
return "nhwc"
|
|
if array.shape[1] == compiled[3]:
|
|
return "nchw"
|
|
raise ValueError(f"binary takes {compiled[3]} channels, fed {list(array.shape)}")
|
|
|
|
|
|
class RknnSession:
|
|
def __init__(self, model_path: Path) -> None:
|
|
self.tpe = settings.rknn_threads or RKNN_NPU_CORES.get(soc_name or "", 1)
|
|
|
|
log.info(f"Loading RKNN model from {model_path} with {self.tpe} threads.")
|
|
self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe)
|
|
log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads sharing its weights.")
|
|
# the shapes the binary was compiled for, among which the runtime routes itself
|
|
batch = self.rknnpool.inputs[0].shape[0]
|
|
declared = orjson.loads(self.rknnpool.custom_string) if self.rknnpool.custom_string else {}
|
|
self.shapes = tuple(Shape(batch, **dims) for dims in declared.get("dims", [{}]))
|
|
self.metadata = {key: value for key, value in declared.items() if key != "dims"}
|
|
self.batches = (batch,)
|
|
|
|
def for_shape(self, shape: Shape) -> RknnSession:
|
|
return self
|
|
|
|
def warm(self) -> None:
|
|
pass # the binary is loaded whole when the session opens
|
|
|
|
def get_inputs(self) -> Sequence[SessionNode]:
|
|
return self.rknnpool.inputs
|
|
|
|
def get_outputs(self) -> Sequence[SessionNode]:
|
|
return self.rknnpool.outputs
|
|
|
|
def get_metadata(self) -> dict[str, str]:
|
|
return self.metadata
|
|
|
|
@property
|
|
def normalizes_input(self) -> bool:
|
|
# the compiler transposes and casts every binary, so only what the exporter declared can say
|
|
return bool(self.rknnpool.custom_string)
|
|
|
|
def run(
|
|
self,
|
|
output_names: list[str] | None,
|
|
input_feed: ModelInput,
|
|
run_options: Any = None,
|
|
) -> list[NDArray[np.float32]]:
|
|
inputs = list(input_feed.values())
|
|
return self.rknnpool.run(inputs, input_layout(self.rknnpool.inputs[0].shape, inputs[0]))
|
|
|
|
|
|
__all__ = ["RknnSession", "RknnNode", "is_available", "soc_name", "model_path"]
|