From 2ce950ac2514d18e2fda0270d07ec1c3f433e747 Mon Sep 17 00:00:00 2001 From: DevTekVE Date: Tue, 7 Jan 2025 07:59:13 +0100 Subject: [PATCH] i think this might work? --- selfdrive/modeld/modeld.py | 127 +++++++++++------------ selfdrive/modeld/runners/model_runner.py | 110 ++++++++++++++++++++ 2 files changed, 171 insertions(+), 66 deletions(-) create mode 100644 selfdrive/modeld/runners/model_runner.py diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 5783f6aa8e..4f01e4e3b9 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,21 +1,13 @@ #!/usr/bin/env python3 -import os from openpilot.system.hardware import TICI +from openpilot.selfdrive.modeld.runners.model_runner import ONNXRunner, TinygradRunner + # -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 -from pathlib import Path from setproctitle import setproctitle from cereal.messaging import PubMaster, SubMaster from msgq.visionipc import VisionIpcClient, VisionStreamType, VisionBuf @@ -31,15 +23,10 @@ from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper from openpilot.selfdrive.modeld.parse_model_outputs import Parser from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState from openpilot.selfdrive.modeld.constants import ModelConstants -from openpilot.selfdrive.modeld.models.commonmodel_pyx import DrivingModelFrame, CLContext - +from openpilot.selfdrive.modeld.models.commonmodel_pyx import DrivingModelFrame_uint8, DrivingModelFrameLegacy as 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' -METADATA_PATH = Path(__file__).parent / 'models/supercombo_metadata.pkl' class FrameMeta: frame_id: int = 0 @@ -59,37 +46,28 @@ class ModelState: def __init__(self, context: CLContext): self.frames = {'input_imgs': DrivingModelFrame(context), 'big_input_imgs': DrivingModelFrame(context)} self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32) + self.full_features_20Hz = np.zeros((ModelConstants.FULL_HISTORY_BUFFER_LEN, ModelConstants.FEATURE_LEN), dtype=np.float32) + self.desire_20Hz = np.zeros((ModelConstants.FULL_HISTORY_BUFFER_LEN + 1, ModelConstants.DESIRE_LEN), dtype=np.float32) + self.is_20hz = False + # Initialize model runner + self.model_runner = TinygradRunner(self.frames) if TICI else ONNXRunner(self.frames) # img buffers are managed in openCL transform code - self.numpy_inputs = { - 'desire': np.zeros((1, (ModelConstants.FULL_HISTORY_BUFFER_LEN+1), ModelConstants.DESIRE_LEN), dtype=np.float32), - 'traffic_convention': np.zeros((1, ModelConstants.TRAFFIC_CONVENTION_LEN), dtype=np.float32), - 'lateral_control_params': np.zeros((1, ModelConstants.LATERAL_CONTROL_PARAMS_LEN), dtype=np.float32), - 'prev_desired_curv': np.zeros((1, (ModelConstants.FULL_HISTORY_BUFFER_LEN+1), ModelConstants.PREV_DESIRED_CURV_LEN), dtype=np.float32), - 'features_buffer': np.zeros((1, ModelConstants.FULL_HISTORY_BUFFER_LEN, ModelConstants.FEATURE_LEN), dtype=np.float32), - } + self.numpy_inputs = {} - with open(METADATA_PATH, 'rb') as f: - model_metadata = pickle.load(f) - self.input_shapes = model_metadata['input_shapes'] + for key, shape in self.model_runner.input_shapes.items(): + if key not in self.frames: # Managed by opencl + self.numpy_inputs[key] = np.zeros(shape, dtype=np.float32) - 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) 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) + net_output_size = self.model_runner.model_metadata['output_shapes']['outputs'][1] + self.output = np.zeros(net_output_size, dtype=np.float32) - 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 + num_elements = self.numpy_inputs['features_buffer'].shape[1] + step_size = int(-100 / num_elements) + self.full_features_20Hz_idxs = np.arange(step_size, step_size * (num_elements + 1), step_size)[::-1] + self.desire_reshape_dims = (self.numpy_inputs['desire'].shape[0], self.numpy_inputs['desire'].shape[1], -1, self.numpy_inputs['desire'].shape[2]) 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: @@ -98,40 +76,56 @@ class ModelState: new_desire = np.where(inputs['desire'] - self.prev_desire > .99, inputs['desire'], 0) self.prev_desire[:] = inputs['desire'] - self.numpy_inputs['desire'][0,:-1] = self.numpy_inputs['desire'][0,1:] - self.numpy_inputs['desire'][0,-1] = new_desire + if self.is_20hz: + self.desire_20Hz[:-1] = self.desire_20Hz[1:] + self.desire_20Hz[-1] = new_desire + self.numpy_inputs['desire'][:] = self.desire_20Hz.reshape(self.desire_reshape_dims).max(axis=2) + else: + self.numpy_inputs['desire'][0,:-1] = self.numpy_inputs['desire'][0,1:] + self.numpy_inputs['desire'][0,-1] = new_desire + + + for key in self.numpy_inputs: + if key in inputs and key not in ['desire']: + self.numpy_inputs[key][:] = inputs[key] - self.numpy_inputs['traffic_convention'][:] = inputs['traffic_convention'] - self.numpy_inputs['lateral_control_params'][:] = inputs['lateral_control_params'] 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]).astype(dtype=np.float32) + # Prepare inputs using the model runner + 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() + # Run model inference + self.output = self.model_runner.run_model() + outputs = self.parser.parse_outputs(self.model_runner.slice_outputs(self.output), self.numpy_inputs.keys()) + + if self.is_20hz: + self.full_features_20Hz[:-1] = self.full_features_20Hz[1:] + self.full_features_20Hz[-1] = outputs['hidden_state'][0, :] + self.numpy_inputs['features_buffer'][:] = self.full_features_20Hz[self.full_features_20Hz_idxs] else: - self.output = self.onnx_cpu_runner.run(None, self.numpy_inputs)[0].flatten() - - outputs = self.parser.parse_outputs(self.slice_outputs(self.output)) - - self.numpy_inputs['features_buffer'][0,:-1] = self.numpy_inputs['features_buffer'][0,1:] - self.numpy_inputs['features_buffer'][0,-1] = outputs['hidden_state'][0, :] + feature_len = outputs['hidden_state'].shape[1] + self.numpy_inputs['features_buffer'][0,:-1] = self.numpy_inputs['features_buffer'][0,1:] + self.numpy_inputs['features_buffer'][0,-1] = outputs['hidden_state'][0, :] + # self.numpy_inputs['features_buffer'][0, :-1] = self.numpy_inputs['features_buffer'][0, 1:] + # self.numpy_inputs['features_buffer'][0, -1, :feature_len] = outputs['hidden_state'][0, :feature_len] - # TODO model only uses last value now - self.numpy_inputs['prev_desired_curv'][0,:-1] = self.numpy_inputs['prev_desired_curv'][0,1:] - self.numpy_inputs['prev_desired_curv'][0,-1,:] = outputs['desired_curvature'][0, :] + if "desired_curvature" in outputs: + input_name_prev = None + + if "prev_desired_curvs" in self.numpy_inputs.keys(): + input_name_prev = 'prev_desired_curvs' + elif "prev_desired_curv" in self.numpy_inputs.keys(): + input_name_prev = 'prev_desired_curv' + + if input_name_prev is not None: + len = outputs['desired_curvature'][0].size + self.numpy_inputs['prev_desired_curv'][0,:-len] = self.numpy_inputs['prev_desired_curv'][0,len:] + self.numpy_inputs['prev_desired_curv'][0,-len,:] = outputs['desired_curvature'][0, :] return outputs @@ -242,7 +236,6 @@ def main(demo=False): is_rhd = sm["driverMonitoringState"].isRHD frame_id = sm["roadCameraState"].frameId v_ego = max(sm["carState"].vEgo, 0.) - lateral_control_params = np.array([v_ego, steer_delay], dtype=np.float32) if sm.updated["liveCalibration"] and sm.seen['roadCameraState'] and sm.seen['deviceState']: device_from_calib_euler = np.array(sm["liveCalibration"].rpyCalib, dtype=np.float32) dc = DEVICE_CAMERAS[(str(sm['deviceState'].deviceType), str(sm['roadCameraState'].sensor))] @@ -273,8 +266,10 @@ def main(demo=False): inputs:dict[str, np.ndarray] = { 'desire': vec_desire, 'traffic_convention': traffic_convention, - 'lateral_control_params': lateral_control_params, - } + } + + if "lateral_control_params" in model.numpy_inputs.keys(): + inputs['lateral_control_params'] = np.array([sm["carState"].vEgo, steer_delay], dtype=np.float32) mt1 = time.perf_counter() model_output = model.run(buf_main, buf_extra, model_transform_main, model_transform_extra, inputs, prepare_only) diff --git a/selfdrive/modeld/runners/model_runner.py b/selfdrive/modeld/runners/model_runner.py new file mode 100644 index 0000000000..3d4c8725e0 --- /dev/null +++ b/selfdrive/modeld/runners/model_runner.py @@ -0,0 +1,110 @@ +import os +from openpilot.system.hardware import TICI + +# +from tinygrad.tensor import Tensor, dtypes +from openpilot.selfdrive.modeld.runners.tinygrad_helpers import qcom_tensor_from_opencl_address +from openpilot.selfdrive.modeld.runners.ort_helpers import make_onnx_cpu_runner, ORT_TYPES_TO_NP_TYPES +import pickle +import numpy as np +from pathlib import Path +from abc import ABC, abstractmethod +from openpilot.selfdrive.modeld.models.commonmodel_pyx import DrivingModelFrame, CLMem + +if TICI: + os.environ['QCOM'] = '1' + +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' +METADATA_PATH = Path(__file__).parent / '../models/supercombo_metadata.pkl' + + +class ModelRunner(ABC): + """Abstract base class for model runners that defines the interface for running ML models.""" + + def __init__(self): + """Initialize the model runner with paths to model and metadata files.""" + 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'] + self.inputs: dict = {} + + @abstractmethod + def prepare_inputs(self, imgs_cl: dict[str, CLMem], numpy_inputs: dict[str, np.ndarray])-> dict: + """Prepare inputs for model inference.""" + + @abstractmethod + def run_model(self): + """Run model inference with prepared inputs.""" + + def slice_outputs(self, model_outputs: np.ndarray) -> dict: + """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, frames: dict[str, DrivingModelFrame] | None = None): + super().__init__() + # Load Tinygrad model + with open(MODEL_PKL_PATH, "rb") as f: + self.model_run = pickle.load(f) + + self.input_to_dtype = {} + self.input_to_device = {} + + for idx, name in enumerate(self.model_run.captured.expected_names): + self.input_to_dtype[name] = self.model_run.captured.expected_st_vars_dtype_device[idx][2] # 2 is the dtype + self.input_to_device[name] = self.model_run.captured.expected_st_vars_dtype_device[idx][3] # 3 is the device + + assert TICI or frames is not None, "TinygradRunner requires frames for non-TICI hardware" + self.frames = frames + self.is_memory_model = None # Use None to indicate that it hasn't been determined yet + + def prepare_inputs(self, imgs_cl: dict[str, CLMem], numpy_inputs: dict[str, np.ndarray]) -> dict: + # Initialize image tensors if not already done + for key in imgs_cl: + if TICI and key not in self.inputs: + self.inputs[key] = qcom_tensor_from_opencl_address(imgs_cl[key].mem_address, self.input_shapes[key], dtype=dtypes.uint8) + elif not TICI: + shape = self.frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.input_shapes[key]) + self.inputs[key] = Tensor(shape, device=self.input_to_device[key], dtype=self.input_to_dtype[key]).realize() + + # Update numpy inputs + for key, value in numpy_inputs.items(): + if key not in imgs_cl: + self.inputs[key] = Tensor(value, device=self.input_to_device[key], dtype=self.input_to_dtype[key]).realize() + + return self.inputs + + def run_model(self): + return self.model_run(**self.inputs).numpy().flatten() + + +class ONNXRunner(ModelRunner): + """ONNX implementation of model runner for non-TICI hardware.""" + + def __init__(self, frames: dict[str, DrivingModelFrame]): + super().__init__() + self.runner = make_onnx_cpu_runner(MODEL_PATH) + self.frames = frames + + self.input_to_nptype = { + model_input.name: ORT_TYPES_TO_NP_TYPES[model_input.type] + for model_input in self.runner.get_inputs() + } + + def prepare_inputs(self, imgs_cl: dict[str, CLMem], numpy_inputs: dict[str, np.ndarray]) -> dict: + self.inputs = numpy_inputs.copy() + for key in imgs_cl: + self.inputs[key] = self.frames[key].buffer_from_cl(imgs_cl[key]).reshape(self.input_shapes[key]).astype(dtype=np.float32) + return self.inputs + + def run_model(self): + return self.runner.run(None, self.inputs)[0].flatten()