diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 49f8c483f0..74b0f80096 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,17 +1,10 @@ #!/usr/bin/env python3 -import os from openpilot.system.hardware import TICI +from openpilot.selfdrive.modeld.runners.model_runner import create_model_runner + # -if TICI: - from tinygrad.tensor import Tensor - from tinygrad.dtype import dtypes - from openpilot.selfdrive.modeld.runners.tinygrad_helpers import qcom_tensor_from_opencl_address - os.environ['QCOM'] = '1' -else: - from openpilot.selfdrive.modeld.runners.ort_helpers import make_onnx_cpu_runner import time -import pickle import numpy as np import cereal.messaging as messaging from cereal import car, log @@ -33,9 +26,7 @@ from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_ from openpilot.selfdrive.modeld.constants import ModelConstants from openpilot.selfdrive.modeld.models.commonmodel_pyx import DrivingModelFrame, CLContext - PROCESS_NAME = "selfdrive.modeld.modeld" -SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') MODEL_PATH = Path(__file__).parent / 'models/supercombo.onnx' MODEL_PKL_PATH = Path(__file__).parent / 'models/supercombo_tinygrad.pkl' @@ -69,27 +60,18 @@ class ModelState: 'features_buffer': np.zeros((1, ModelConstants.HISTORY_BUFFER_LEN, ModelConstants.FEATURE_LEN), dtype=np.float32), } - with open(METADATA_PATH, 'rb') as f: - model_metadata = pickle.load(f) - self.input_shapes = model_metadata['input_shapes'] - - self.output_slices = model_metadata['output_slices'] - net_output_size = model_metadata['output_shapes']['outputs'][1] - self.output = np.zeros(net_output_size, dtype=np.float32) + # Initialize model runner + self.model_runner = create_model_runner( + model_path=MODEL_PATH, + metadata_path=METADATA_PATH, + frames=self.frames, + tinygrad_path=MODEL_PKL_PATH, + is_tici=TICI + ) self.parser = Parser() - if TICI: - self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} - with open(MODEL_PKL_PATH, "rb") as f: - self.model_run = pickle.load(f) - else: - self.onnx_cpu_runner = make_onnx_cpu_runner(MODEL_PATH) - - def slice_outputs(self, model_outputs: np.ndarray) -> dict[str, np.ndarray]: - parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in self.output_slices.items()} - if SEND_RAW_PRED: - parsed_model_outputs['raw_pred'] = model_outputs.copy() - return parsed_model_outputs + net_output_size = self.model_runner.model_metadata['output_shapes']['outputs'][1] + self.output = np.zeros(net_output_size, dtype=np.float32) def run(self, buf: VisionBuf, wbuf: VisionBuf, transform: np.ndarray, transform_wide: np.ndarray, inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None: @@ -106,24 +88,15 @@ class ModelState: imgs_cl = {'input_imgs': self.frames['input_imgs'].prepare(buf, transform.flatten()), 'big_input_imgs': self.frames['big_input_imgs'].prepare(wbuf, transform_wide.flatten())} - if TICI: - # The imgs tensors are backed by opencl memory, only need init once - for key in imgs_cl: - if key not in self.tensor_inputs: - self.tensor_inputs[key] = qcom_tensor_from_opencl_address(imgs_cl[key].mem_address, self.input_shapes[key], dtype=dtypes.uint8) - else: - for key in imgs_cl: - self.numpy_inputs[key] = self.frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.input_shapes[key]) + # Prepare inputs using the model runner + prepared_inputs = self.model_runner.prepare_inputs(imgs_cl, self.numpy_inputs) if prepare_only: return None - if TICI: - self.output = self.model_run(**self.tensor_inputs).numpy().flatten() - else: - self.output = self.onnx_cpu_runner.run(None, self.numpy_inputs)[0].flatten() - - outputs = self.parser.parse_outputs(self.slice_outputs(self.output)) + # Run model inference + self.output = self.model_runner.run_model(prepared_inputs) + outputs = self.parser.parse_outputs(self.model_runner.slice_outputs(self.output)) self.full_features_20Hz[:-1] = self.full_features_20Hz[1:] self.full_features_20Hz[-1] = outputs['hidden_state'][0, :] @@ -190,7 +163,6 @@ def main(demo=False): meta_main = FrameMeta() meta_extra = FrameMeta() - if demo: CP = get_demo_car_params() else: diff --git a/selfdrive/modeld/runners/__init__.py b/selfdrive/modeld/runners/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/selfdrive/modeld/runners/model_runner.py b/selfdrive/modeld/runners/model_runner.py new file mode 100644 index 0000000000..74b03c791e --- /dev/null +++ b/selfdrive/modeld/runners/model_runner.py @@ -0,0 +1,111 @@ +import os +from openpilot.system.hardware import TICI + +# +if TICI: + from tinygrad.tensor import Tensor + from tinygrad.dtype import dtypes + from openpilot.selfdrive.modeld.runners.tinygrad_helpers import qcom_tensor_from_opencl_address + + os.environ['QCOM'] = '1' +else: + from openpilot.selfdrive.modeld.runners.ort_helpers import make_onnx_cpu_runner +import pickle +import numpy as np +from pathlib import Path +from abc import ABC, abstractmethod +from openpilot.selfdrive.modeld.models.commonmodel_pyx import DrivingModelFrame, CLContext + +SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') + + +class ModelRunner(ABC): + """Abstract base class for model runners that defines the interface for running ML models.""" + + def __init__(self, model_path: Path, metadata_path: Path, frames: dict[str, DrivingModelFrame]): + """Initialize the model runner with paths to model and metadata files.""" + self.model_path = model_path + self.frames = frames + with open(metadata_path, 'rb') as f: + self.model_metadata = pickle.load(f) + self.input_shapes = self.model_metadata['input_shapes'] + self.output_slices = self.model_metadata['output_slices'] + + @abstractmethod + def prepare_inputs(self, imgs_cl: dict[str, any], numpy_inputs: dict[str, np.ndarray]) -> dict[str, any]: + """Prepare inputs for model inference.""" + pass + + @abstractmethod + def run_model(self, inputs: dict[str, any]) -> np.ndarray: + """Run model inference with prepared inputs.""" + pass + + def slice_outputs(self, model_outputs: np.ndarray) -> dict[str, np.ndarray]: + """Slice model outputs according to metadata configuration.""" + parsed_outputs = {k: model_outputs[np.newaxis, v] for k, v in self.output_slices.items()} + if SEND_RAW_PRED: + parsed_outputs['raw_pred'] = model_outputs.copy() + return parsed_outputs + + +class TinyGradRunner(ModelRunner): + """TinyGrad implementation of model runner for TICI hardware.""" + + def __init__(self, model_path: Path, metadata_path: Path, frames: dict[str, DrivingModelFrame]): + super().__init__(model_path, metadata_path, frames) + # Load TinyGrad model + with open(model_path, "rb") as f: + self.model_run = pickle.load(f) + self.tensor_inputs = {} + + def prepare_inputs(self, imgs_cl: dict[str, any], numpy_inputs: dict[str, np.ndarray]) -> dict[str, any]: + # Initialize image tensors if not already done + for key in imgs_cl: + if key not in self.tensor_inputs: + self.tensor_inputs[key] = qcom_tensor_from_opencl_address( + imgs_cl[key].mem_address, + self.input_shapes[key], + dtype=dtypes.uint8 + ) + + # Update numpy inputs + for k, v in numpy_inputs.items(): + if k not in self.tensor_inputs: + self.tensor_inputs[k] = Tensor(v, device='NPY').realize() + else: + self.tensor_inputs[k].assign(v) + + return self.tensor_inputs + + def run_model(self, inputs: dict[str, any]) -> np.ndarray: + return self.model_run(**inputs).flatten() + + +class ONNXRunner(ModelRunner): + """ONNX implementation of model runner for non-TICI hardware.""" + + def __init__(self, model_path: Path, metadata_path: Path, frames: dict[str, DrivingModelFrame]): + super().__init__(model_path, metadata_path, frames) + self.runner = make_onnx_cpu_runner(model_path) + + def prepare_inputs(self, imgs_cl: dict[str, any], numpy_inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + for key in imgs_cl: + numpy_inputs[key] = self.frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.input_shapes[key]) + return numpy_inputs + + def run_model(self, inputs: dict[str, any]) -> np.ndarray: + return self.runner.run(None, inputs)[0].flatten() + + +def create_model_runner( + model_path: Path, + metadata_path: Path, + frames: dict[str, DrivingModelFrame], + tinygrad_path: Path | None = None, + is_tici: bool = False +) -> ModelRunner: + """Factory function to create appropriate model runner based on hardware.""" + if is_tici: + return TinyGradRunner(tinygrad_path or model_path, metadata_path, frames) + return ONNXRunner(model_path, metadata_path, frames)