191 lines
7.5 KiB
Python
191 lines
7.5 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pickle
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
from tinygrad.dtype import dtypes
|
|
from tinygrad.tensor import Tensor
|
|
|
|
from iqpilot.selfdrive.iqmodeld.models.runners.model_runner import (
|
|
CLMemDict,
|
|
CUSTOM_MODEL_PATH,
|
|
FrameDict,
|
|
ModelType,
|
|
NumpyDict,
|
|
ShapeDict,
|
|
SliceDict,
|
|
)
|
|
from iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
|
from iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.model_types import (
|
|
OffPolicyTinygrad,
|
|
OnPolicyTinygrad,
|
|
PolicyTinygrad,
|
|
SupercomboTinygrad,
|
|
VisionTinygrad,
|
|
)
|
|
from iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
|
|
from iqpilot.selfdrive.iqmodeld.config import ModelConstants
|
|
from iqpilot.selfdrive.iqmodeld.runtime.tinygrad import qcom_tensor_from_opencl_address
|
|
from iqpilot.system.hardware import TICI
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _TensorShapePlan:
|
|
dtype: object
|
|
device: str
|
|
|
|
|
|
def _artifact_path(filename: str) -> str:
|
|
return f"{CUSTOM_MODEL_PATH}/{filename}"
|
|
|
|
|
|
def _load_program_blob(filename: str):
|
|
with open(_artifact_path(filename), "rb") as artifact:
|
|
try:
|
|
return pickle.load(artifact)
|
|
except FileNotFoundError as exc:
|
|
assert "/dev/kgsl-3d0" not in str(exc), "Model was built on C3 or C3X, but is being loaded on PC"
|
|
raise
|
|
|
|
|
|
def _compile_input_plan(captured) -> dict[str, _TensorShapePlan]:
|
|
plan: dict[str, _TensorShapePlan] = {}
|
|
for name, info in zip(captured.expected_names, captured.expected_input_info, strict=True):
|
|
plan[name] = _TensorShapePlan(dtype=info[2], device=info[3])
|
|
return plan
|
|
|
|
|
|
def _merge_step_outputs(output_groups: list[NumpyDict]) -> NumpyDict:
|
|
stitched: NumpyDict = {}
|
|
for payload in output_groups:
|
|
stitched.update(payload)
|
|
if "planplus" in stitched and "plan" in stitched:
|
|
stitched["plan"] = stitched["plan"] + stitched["planplus"]
|
|
return stitched
|
|
|
|
|
|
class TinygradRunner(ModelRunner, SupercomboTinygrad, PolicyTinygrad, VisionTinygrad, OffPolicyTinygrad, OnPolicyTinygrad):
|
|
def __init__(self, model_type: int = ModelType.supercombo):
|
|
ModelRunner.__init__(self)
|
|
for initializer in (SupercomboTinygrad, PolicyTinygrad, VisionTinygrad, OffPolicyTinygrad, OnPolicyTinygrad):
|
|
initializer.__init__(self)
|
|
|
|
self._constants = ModelConstants
|
|
self._model_data = self.models.get(model_type)
|
|
if self._model_data is None or self._model_data.model is None:
|
|
raise ValueError(f"Model data for type {model_type} not available.")
|
|
|
|
asset_name = self._model_data.model.artifact.fileName
|
|
assert asset_name.endswith("_tinygrad.pkl"), f"Invalid model file {asset_name} for TinygradRunner"
|
|
|
|
self.model_run = _load_program_blob(asset_name)
|
|
self._input_plan = _compile_input_plan(self.model_run.captured)
|
|
for name, spec in self._input_plan.items():
|
|
if "img" in name and spec.dtype is not dtypes.uint8:
|
|
raise ValueError(f"{asset_name}: image input {name} expects {spec.dtype}, incompatible with uint8 warp buffer")
|
|
self.input_to_dtype = {name: spec.dtype for name, spec in self._input_plan.items()}
|
|
self.input_to_device = {name: spec.device for name, spec in self._input_plan.items()}
|
|
|
|
@property
|
|
def vision_input_names(self) -> list[str]:
|
|
return [stream_name for stream_name in self.input_shapes if "img" in stream_name]
|
|
|
|
def _attach_vision_tensor(self, stream_name: str, frame_buffers: CLMemDict, frame_views: FrameDict) -> None:
|
|
spec = self._input_plan[stream_name]
|
|
frame_buffer = frame_buffers[stream_name]
|
|
if TICI:
|
|
self.inputs[stream_name] = qcom_tensor_from_opencl_address(frame_buffer.mem_address,
|
|
self.input_shapes[stream_name],
|
|
dtype=spec.dtype)
|
|
return
|
|
|
|
mirrored = frame_views[stream_name].as_numpy(frame_buffer).reshape(self.input_shapes[stream_name])
|
|
self.inputs[stream_name] = Tensor(mirrored, device=spec.device, dtype=spec.dtype).realize()
|
|
|
|
def _attach_state_tensor(self, tensor_name: str, tensor_value: np.ndarray) -> None:
|
|
spec = self._input_plan[tensor_name]
|
|
self.inputs[tensor_name] = Tensor(tensor_value, device=spec.device, dtype=spec.dtype).realize()
|
|
|
|
def prepare_vision_inputs(self, imgs_cl: CLMemDict, frames: FrameDict):
|
|
for stream_name in imgs_cl:
|
|
if stream_name not in self.inputs or not TICI:
|
|
self._attach_vision_tensor(stream_name, imgs_cl, frames)
|
|
|
|
def prepare_policy_inputs(self, numpy_inputs: NumpyDict):
|
|
for tensor_name, tensor_value in numpy_inputs.items():
|
|
self._attach_state_tensor(tensor_name, tensor_value)
|
|
|
|
def prepare_inputs(self, imgs_cl: CLMemDict, numpy_inputs: NumpyDict, frames: FrameDict) -> dict:
|
|
self.prepare_vision_inputs(imgs_cl, frames)
|
|
self.prepare_policy_inputs(numpy_inputs)
|
|
return self.inputs
|
|
|
|
def _parse_outputs(self, model_outputs: np.ndarray) -> NumpyDict:
|
|
if self._model_data is None:
|
|
raise ValueError("Model data is not available. Ensure the model is loaded correctly.")
|
|
return self.parser_method_dict[self._model_data.model.type.raw](model_outputs)
|
|
|
|
def _run_model(self) -> NumpyDict:
|
|
raw_output = self.model_run(**self.inputs).numpy().reshape(-1)
|
|
return self._parse_outputs(raw_output)
|
|
|
|
|
|
class TinygradSplitRunner(ModelRunner):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.is_20hz_3d = True
|
|
self._constants = SplitModelConstants
|
|
self.vision_runner = TinygradRunner(ModelType.vision)
|
|
self.policy_runner = TinygradRunner(ModelType.policy) if self.models.get(ModelType.policy) else None
|
|
self.off_policy_runner = TinygradRunner(ModelType.offPolicy) if self.models.get(ModelType.offPolicy) else None
|
|
self.on_policy_runner = TinygradRunner(ModelType.onPolicy) if self.models.get(ModelType.onPolicy) else None
|
|
|
|
def _policy_units(self) -> list[TinygradRunner]:
|
|
return [runner for runner in (self.policy_runner, self.off_policy_runner, self.on_policy_runner) if runner is not None]
|
|
|
|
def run_vision(self) -> NumpyDict:
|
|
return self.vision_runner.run_model()
|
|
|
|
def run_policy(self) -> NumpyDict:
|
|
return _merge_step_outputs([runner.run_model() for runner in self._policy_units()])
|
|
|
|
def refresh_policy_features(self, features_buffer: np.ndarray) -> None:
|
|
for runner in self._policy_units():
|
|
if "features_buffer" in runner._input_plan:
|
|
runner._attach_state_tensor("features_buffer", features_buffer)
|
|
|
|
def _run_model(self) -> NumpyDict:
|
|
return _merge_step_outputs([self.run_vision(), self.run_policy()])
|
|
|
|
@property
|
|
def vision_input_names(self) -> list[str]:
|
|
return list(self.vision_runner.vision_input_names)
|
|
|
|
@property
|
|
def input_shapes(self) -> ShapeDict:
|
|
composite: ShapeDict = dict(self.vision_runner.input_shapes)
|
|
for runner in self._policy_units():
|
|
composite.update(runner.input_shapes)
|
|
return composite
|
|
|
|
@property
|
|
def output_slices(self) -> SliceDict:
|
|
composite: SliceDict = dict(self.vision_runner.output_slices)
|
|
for runner in self._policy_units():
|
|
composite.update(runner.output_slices)
|
|
return composite
|
|
|
|
def prepare_inputs(self, imgs_cl: CLMemDict, numpy_inputs: NumpyDict, frames: FrameDict) -> dict:
|
|
self.vision_runner.prepare_vision_inputs(imgs_cl, frames)
|
|
assembled_inputs = dict(self.vision_runner.inputs)
|
|
for runner in self._policy_units():
|
|
runner.prepare_policy_inputs(numpy_inputs)
|
|
assembled_inputs.update(runner.inputs)
|
|
self.inputs = assembled_inputs
|
|
return assembled_inputs
|