diff --git a/openpilot/sunnypilot/modeld_v2/__init__.py b/openpilot/sunnypilot/modeld_v2/__init__.py index 639622e827..e69de29bb2 100644 --- a/openpilot/sunnypilot/modeld_v2/__init__.py +++ b/openpilot/sunnypilot/modeld_v2/__init__.py @@ -1,5 +0,0 @@ -from pathlib import Path - -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' diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld.py b/openpilot/sunnypilot/modeld_v2/compile_modeld.py index a274097563..1e88769928 100755 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld.py @@ -32,7 +32,7 @@ def _patch_tinygrad_fetch_fw(): helpers.fetch_fw = fetch_fw _patch_tinygrad_fetch_fw() -from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare, sample_desire, sample_skip, shift_and_sample +import openpilot.selfdrive.modeld.compile_modeld as stock from tinygrad import dtypes from tinygrad.device import Device from tinygrad.engine.jit import TinyJit @@ -41,7 +41,7 @@ from tinygrad.tensor import Tensor MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy') WARP_INPUTS = ['tfm', 'big_tfm'] POLICY_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs'] - +nv12_copy_size = stock.nv12_copy_size def _detect_desire_key(shapes: dict) -> str | None: return next((key for key in shapes if key.startswith('desire')), None) @@ -138,7 +138,7 @@ def make_supercombo_input_queues(input_shapes: dict, frame_skip: int, return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True) -def make_random_images(keys, shape, device): +def make_random_images(keys, shape, device, rng=None): return {k: Tensor.randint(shape, low=0, high=256, dtype=dtypes.uint8, device=device).realize() for k in keys} @@ -151,25 +151,9 @@ def make_warp_queues(device=Device.DEFAULT): return queues, npy -def make_warp(nv12: NV12Frame, model_w: int, model_h: int): - frame_prepare = make_frame_prepare(nv12, model_w, model_h) - - def warp(tfm, big_tfm, frame, big_frame): - tfm = tfm.to(Device.DEFAULT) - big_tfm = big_tfm.to(Device.DEFAULT) - frame = frame.to(Device.DEFAULT) - big_frame = big_frame.to(Device.DEFAULT) - Tensor.realize(tfm, big_tfm, frame, big_frame) - - warped_frame = frame_prepare(frame, tfm).unsqueeze(0) - warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0) - return Tensor.cat(warped_frame, warped_big_frame) - return warp - - def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, frame_skip: int, input_shapes: dict): - sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) - sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) + sample_skip_fn = partial(stock.sample_skip, frame_skip=frame_skip) + sample_desire_fn = partial(stock.sample_desire, frame_skip=frame_skip) desire_key = _detect_desire_key(input_shapes) road_key, wide_key = _detect_vision_keys(input_shapes) @@ -186,14 +170,14 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, warped_dev = warped.to(Device.DEFAULT) Tensor.realize(packed_npy_inputs_dev, warped_dev) - img = shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn) - big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn) + img = stock.shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn) + big_img = stock.shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn) unpacked_tensors = [tensor.reshape(shape) for tensor, shape in zip(packed_npy_inputs_dev.split(npy_sizes), npy_shapes.values(), strict=True)] unpacked_dict = dict(zip(npy_shapes.keys(), unpacked_tensors, strict=True)) desire_dev = unpacked_dict['desire'] - desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn) + desire_buf = stock.shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn) inputs = {desire_key: desire_buf} for key, tensor_val in unpacked_dict.items(): @@ -202,13 +186,13 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, if 'prev_feat' in unpacked_dict: prev_feat_dev = unpacked_dict['prev_feat'] - inputs['features_buffer'] = shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).reshape(input_shapes['features_buffer']) + inputs['features_buffer'] = stock.shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).reshape(input_shapes['features_buffer']) if vision_runner: vision_out_cast = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast('float32').realize() if 'features_buffer' not in inputs: new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0) - inputs['features_buffer'] = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() + inputs['features_buffer'] = stock.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() policy_outs = [next(iter(pol_runner(inputs).values())).cast('float32').realize() for pol_runner in policy_runners] return (vision_out_cast, *policy_outs) if len(policy_outs) > 1 else (vision_out_cast, policy_outs[0]) @@ -219,27 +203,28 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice, policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32').realize() if 'features_buffer' not in inputs and features_slice is not None: new_feat = policy_out[:, features_slice].reshape(1, -1).unsqueeze(0) - shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() + stock.shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() return policy_out return run_policy -def compile_jit(jit, make_random_inputs, input_keys, make_queues): +def compile_jit(jit, input_keys, make_queues, make_random_inputs=None, benchmark_runs: int = 1): SEED = 42 - def random_inputs_run(fn, seed, test_val=None, test_buffers=None, expect_match=True): - input_queues, npy = make_queues(Device.DEFAULT) + def random_inputs_run(fn, seed, n_runs, test_val=None, test_buffers=None, expect_match=True): + queues_res = make_queues(Device.DEFAULT) + input_queues, npy = queues_res[0], queues_res[1] + frame_views = queues_res[2] if len(queues_res) > 2 else {} rng = np.random.default_rng(seed) Tensor.manual_seed(seed) - testing = test_val is not None or test_buffers is not None - n_runs = 1 if testing else 3 - for i in range(n_runs): for v in npy.values(): v[:] = rng.standard_normal(v.shape).astype(v.dtype) + for v in frame_views.values(): + v[:] = rng.integers(0, 256, size=v.shape, dtype=np.uint8) Device.default.synchronize() - random_inputs = make_random_inputs() + random_inputs = make_random_inputs(rng=rng) if make_random_inputs is not None else {} st = time.perf_counter() outs = fn(**{k: input_queues[k] for k in input_keys if k in input_queues}, **random_inputs) mt = time.perf_counter() @@ -260,14 +245,15 @@ def compile_jit(jit, make_random_inputs, input_keys, make_queues): return val, buffers print('capture + replay') - test_val, test_buffers = random_inputs_run(jit, SEED) - print('pickle round trip') + test_val, test_buffers = random_inputs_run(jit, SEED, 3) + print(f'pickle round trip ({benchmark_runs} runs per seed)') with tempfile.TemporaryFile(dir=".") as f: dump_oob(jit, f) f.seek(0) - deserialized_jit = load_oob(f) - random_inputs_run(deserialized_jit, SEED, test_val=test_val, test_buffers=test_buffers) - return deserialized_jit + loaded_jit = load_oob(f) + random_inputs_run(loaded_jit, SEED, benchmark_runs, test_val, test_buffers, expect_match=True) + random_inputs_run(loaded_jit, SEED+1, benchmark_runs, test_val, test_buffers, expect_match=False) + return jit def _parse_size(size_str: str) -> tuple[int, int]: @@ -317,6 +303,7 @@ if __name__ == "__main__": parser.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH') parser.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True) parser.add_argument('--frame-skip', type=int, default=None, help='frame skip value (auto-derived if not provided)') + parser.add_argument('--benchmark-runs', type=int, default=1, help='benchmark runs') parser.add_argument('--output', required=True) parser.add_argument('--vision-onnx', help='vision ONNX (for split models)') @@ -335,50 +322,64 @@ if __name__ == "__main__": args.on_policy_onnx = read_file_chunked_to_disk(args.on_policy_onnx) args.supercombo_onnx = read_file_chunked_to_disk(args.supercombo_onnx) - vision_runner = OnnxRunner(args.vision_onnx) if args.vision_onnx else None - - if args.model_type == 'vision_policy': - assert vision_runner and args.policy_onnx - policy_runners = [OnnxRunner(args.policy_onnx)] - output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx), 'policy': make_metadata_dict(args.policy_onnx)} - elif args.model_type == 'supercombo': + if args.model_type == 'supercombo': assert args.supercombo_onnx - policy_runners = [OnnxRunner(args.supercombo_onnx)] - output_data['metadata'] = {'model': make_metadata_dict(args.supercombo_onnx)} - elif args.model_type == 'vision_multi_policy': - assert vision_runner - policy_runners, policy_names = _load_policy_runners(args) - output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx)} - for name in policy_names: - runner_arg = getattr(args, f"{name}_onnx") - output_data['metadata'][name] = make_metadata_dict(runner_arg) + model_metadata = make_metadata_dict(args.supercombo_onnx) + output_data['metadata'] = {'model': model_metadata, **model_metadata} + output_data['input_devices'] = {'model': Device.DEFAULT} + output_data['run_model'] = {} + derived_frame_skip = args.frame_skip or derive_frame_skip({}, model_metadata['input_shapes']) + model_runner = OnnxRunner(args.supercombo_onnx) + run_policy = stock.make_run_policy(model_runner, model_metadata, derived_frame_skip) + for cam_w, cam_h in args.camera_resolutions: + print(f"Compiling unified run_model JIT for {cam_w}x{cam_h}...") + nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height) + make_model_queues = partial(stock.make_input_queues, model_metadata['input_shapes'], derived_frame_skip, + frame_copy_size=frame_copy_size) + warp = stock.make_warp(nv12, model_w, model_h) + run_model_jit = TinyJit(stock.make_run_model(warp, run_policy, model_metadata, frame_copy_size), prune=True) + output_data['run_model'][(cam_w, cam_h)] = compile_jit(run_model_jit, stock.MODELD_INPUTS, make_model_queues, benchmark_runs=args.benchmark_runs) + else: + vision_runner = OnnxRunner(args.vision_onnx) if args.vision_onnx else None + if args.model_type == 'vision_policy': + assert vision_runner and args.policy_onnx + policy_runners = [OnnxRunner(args.policy_onnx)] + output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx), 'policy': make_metadata_dict(args.policy_onnx)} + elif args.model_type == 'vision_multi_policy': + assert vision_runner + policy_runners, policy_names = _load_policy_runners(args) + output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx)} + for name in policy_names: + runner_arg = getattr(args, f"{name}_onnx") + output_data['metadata'][name] = make_metadata_dict(runner_arg) - policy_keys = [key for key in output_data['metadata'].keys() if key != 'vision'] - first_policy_meta = output_data['metadata'][policy_keys[0]] if policy_keys else {} - vision_meta = output_data['metadata'].get('vision', {}) + policy_keys = [key for key in output_data['metadata'].keys() if key != 'vision'] + first_policy_meta = output_data['metadata'][policy_keys[0]] if policy_keys else {} + vision_meta = output_data['metadata'].get('vision', {}) - derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {})) - all_shapes = {key: value for meta in output_data['metadata'].values() for key, value in meta['input_shapes'].items()} - feat_meta = output_data['metadata'].get('vision') or output_data['metadata'].get('model') or output_data['metadata'].get('policy') - assert feat_meta is not None - features_slice = feat_meta['output_slices']['hidden_state'] - is_supercombo = vision_runner is None + derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {})) + all_shapes = {key: value for meta in output_data['metadata'].values() for key, value in meta['input_shapes'].items()} + feat_meta = output_data['metadata'].get('vision') or output_data['metadata'].get('policy') + assert feat_meta is not None + features_slice = feat_meta['output_slices']['hidden_state'] - print(f"Compiling run_policy JIT (model_size={model_w}x{model_h}, frame_skip={derived_frame_skip})...") - run_policy_func = make_run_policy(vision_runner, policy_runners, features_slice, derived_frame_skip, all_shapes) - run_policy_jit = TinyJit(run_policy_func, prune=True) - make_policy_queues = partial(generate_queues_and_npy, all_shapes, derived_frame_skip, is_supercombo=is_supercombo) - make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, model_h // 2, model_w // 2), device=Device.DEFAULT) - output_data['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS, make_policy_queues) + print(f"Compiling run_policy JIT (model_size={model_w}x{model_h}, frame_skip={derived_frame_skip})...") + run_policy_func = make_run_policy(vision_runner, policy_runners, features_slice, derived_frame_skip, all_shapes) + run_policy_jit = TinyJit(run_policy_func, prune=True) + make_policy_queues = partial(generate_queues_and_npy, all_shapes, derived_frame_skip, is_supercombo=False) + make_random_model_inputs = partial(make_random_images, keys=['warped'], shape=(2, 6, model_h // 2, model_w // 2), device=Device.DEFAULT) + output_data['run_policy'] = compile_jit(run_policy_jit, POLICY_INPUTS, make_policy_queues, make_random_inputs=make_random_model_inputs) - for cam_w, cam_h in args.camera_resolutions: - print(f"Compiling warp JIT for {cam_w}x{cam_h}...") - nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) - make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=Device.DEFAULT) - warp = TinyJit(make_warp(nv12, model_w, model_h), prune=True) - output_data[(cam_w, cam_h)] = compile_jit(warp, make_random_warp_inputs, WARP_INPUTS, make_warp_queues) + for cam_w, cam_h in args.camera_resolutions: + print(f"Compiling warp JIT for {cam_w}x{cam_h}...") + nv12 = stock.NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + frame_copy_size = stock.nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height) + make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=frame_copy_size, device=Device.DEFAULT) + warp = TinyJit(stock.make_warp(nv12, model_w, model_h), prune=True) + output_data[(cam_w, cam_h)] = compile_jit(warp, WARP_INPUTS, make_warp_queues, make_random_inputs=make_random_warp_inputs) - output_data['metadata']['warp_dev'] = Device.DEFAULT + output_data['metadata']['warp_dev'] = Device.DEFAULT with open(args.output, "wb") as file: dump_oob(output_data, file) diff --git a/openpilot/sunnypilot/modeld_v2/constants.py b/openpilot/sunnypilot/modeld_v2/constants.py index cf5157591e..c51996b537 100644 --- a/openpilot/sunnypilot/modeld_v2/constants.py +++ b/openpilot/sunnypilot/modeld_v2/constants.py @@ -14,6 +14,8 @@ class ModelConstants: # model inputs constants MODEL_FREQ = 20 + MODEL_RUN_FREQ = 20 + MODEL_CONTEXT_FREQ = 5 FEATURE_LEN = 512 FULL_HISTORY_BUFFER_LEN = 99 DESIRE_LEN = 8 @@ -35,6 +37,7 @@ class ModelConstants: LANE_LINES_WIDTH = 2 ROAD_EDGES_WIDTH = 2 PLAN_WIDTH = 15 + ACTION_WIDTH = 2 DESIRE_PRED_WIDTH = 8 LAT_PLANNER_SOLUTION_WIDTH = 4 DESIRED_CURV_WIDTH = 1 diff --git a/openpilot/sunnypilot/modeld_v2/meta_helper.py b/openpilot/sunnypilot/modeld_v2/meta_helper.py index bcf73adc6f..bf4116bb9b 100644 --- a/openpilot/sunnypilot/modeld_v2/meta_helper.py +++ b/openpilot/sunnypilot/modeld_v2/meta_helper.py @@ -1,26 +1,9 @@ from openpilot.sunnypilot.modeld_v2.constants import Meta -from openpilot.cereal import custom from openpilot.sunnypilot.modeld_v2.meta_20hz import Meta20hz from openpilot.sunnypilot.models.helpers import get_active_bundle -ModelBundle = custom.ModelManagerSP.ModelBundle - def load_meta_constants(): - """ - Determines and loads the appropriate meta model class based on the metadata provided. The function checks - specific keys and conditions within the provided metadata dictionary to identify the corresponding meta - model class to return. - - :param model_metadata: Dictionary containing metadata about the model. It includes - details such as input shapes, output slices, and other configurations for identifying - metadata-dependent meta model classes. - :type model_metadata: dict - :return: The appropriate meta model class (Meta, MetaSimPose, or MetaTombRaider) - based on the conditions and metadata provided. - :rtype: type - """ if (bundle := get_active_bundle()) and bundle.is20hz: return Meta20hz - - return Meta # Default + return Meta diff --git a/openpilot/sunnypilot/modeld_v2/modeld.py b/openpilot/sunnypilot/modeld_v2/modeld.py index 2ab96934ef..683def6588 100755 --- a/openpilot/sunnypilot/modeld_v2/modeld.py +++ b/openpilot/sunnypilot/modeld_v2/modeld.py @@ -38,11 +38,18 @@ from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value from openpilot.selfdrive.modeld.modeld import ChestnutState +from openpilot.selfdrive.modeld.compile_modeld import ( + MODELD_INPUTS, + make_input_queues as make_stock_input_queues, +) from openpilot.sunnypilot.modeld_v2.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState, get_curvature_from_output -from openpilot.sunnypilot.modeld_v2.constants import Plan +from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser +from openpilot.sunnypilot.modeld_v2.constants import ModelConstants, Plan from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper -from openpilot.sunnypilot.modeld_v2.compile_modeld import derive_frame_skip, make_split_input_queues, make_supercombo_input_queues, WARP_INPUTS, POLICY_INPUTS +from openpilot.sunnypilot.modeld_v2.compile_modeld import (derive_frame_skip, make_split_input_queues, + make_supercombo_input_queues, nv12_copy_size, + WARP_INPUTS, POLICY_INPUTS) from openpilot.sunnypilot.livedelay.helpers import get_lat_delay from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase from openpilot.sunnypilot.models.helpers import get_active_bundle @@ -112,29 +119,39 @@ class ModelState(ModelStateBase): jits = load_oob(open_file_chunked(pkl_path)) metadata = jits['metadata'] - self.WARP_DEV = metadata.get('warp_dev', 'QCOM' if COMMA_HARDWARE else 'CPU') - self.DEV = 'AMD' if self.chestnut else ('QCOM' if COMMA_HARDWARE else 'CPU') + self.WARP_DEV = metadata.get('warp_dev', 'QCOM') if COMMA_HARDWARE else 'CPU' + self.DEV = ('AMD' if self.chestnut else 'QCOM') if COMMA_HARDWARE else 'CPU' self.QUEUE_DEV = self.DEV - self.run_policy = jits['run_policy'] - self.warp = jits[(cam_w, cam_h)] + self.is_run_model = 'run_model' in jits - if 'model' in metadata: - model_metadata = metadata['model'] + nv12_info = get_nv12_info(cam_w, cam_h) + self.frame_copy_size = nv12_copy_size(*nv12_info[:3]) + self.full_frames: dict = {} + self._blob_cache: dict = {} + self.frame_buffers: dict = {} + + if self.is_run_model or 'model' in metadata: + model_metadata = metadata.get('model', metadata) + self.input_shapes = model_metadata['input_shapes'] self.vision_output_slices = model_metadata['output_slices'] self.policy_output_slices = {} self._policy_slices_list = [] self._combined_model_type = 'supercombo' - self._vision_input_names = [key for key in model_metadata['input_shapes'] if 'img' in key] - frame_skip = derive_frame_skip({}, model_metadata['input_shapes']) - self.input_queues, self.numpy_inputs = make_supercombo_input_queues(model_metadata['input_shapes'], - frame_skip, device=self.QUEUE_DEV) - else: - vision_metadata = metadata['vision'] - policy_keys = [k for k in metadata if k != 'vision'] - if policy_keys == ['policy']: - self._combined_model_type = 'split' + self._vision_input_names = [key for key in self.input_shapes if 'img' in key] + self.frame_skip = derive_frame_skip({}, self.input_shapes) + if self.is_run_model: + self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues( + self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size) + self.frame_views, self.npy = self.frame_buffers, self.numpy_inputs + self.run_model, self.run_policy, self.warp = jits['run_model'][(cam_w, cam_h)], None, None else: - self._combined_model_type = 'multi_policy' + self.input_queues, self.numpy_inputs = make_supercombo_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV) + self.run_model, self.run_policy, self.warp = None, jits['run_policy'], jits[(cam_w, cam_h)] + else: + self.run_model, self.run_policy, self.warp = None, jits['run_policy'], jits[(cam_w, cam_h)] + vision_metadata = metadata['vision'] + policy_keys = [k for k in metadata if k not in ('vision', 'warp_dev')] + self._combined_model_type = 'split' if policy_keys == ['policy'] else 'multi_policy' self.vision_output_slices = vision_metadata['output_slices'] self._policy_keys = policy_keys self._policy_slices_list = [metadata[k]['output_slices'] for k in policy_keys] @@ -150,50 +167,39 @@ class ModelState(ModelStateBase): self._desire_key = next(key for key in self.numpy_inputs if key.startswith('desire')) self._road_key = next(key for key in self._vision_input_names if 'big' not in key) self._wide_key = next(key for key in self._vision_input_names if 'big' in key) + self.frame_buf_params = dict.fromkeys(self._vision_input_names, nv12_info) is_20hz = bundle.is20hz if bundle else self._combined_model_type in ('split', 'multi_policy') if is_20hz: from openpilot.sunnypilot.models.split_model_constants import SplitModelConstants self.constants = SplitModelConstants() else: - from openpilot.sunnypilot.modeld_v2.constants import ModelConstants self.constants = ModelConstants() - if self._combined_model_type != 'supercombo': - from openpilot.sunnypilot.modeld_v2.parse_model_outputs_split import Parser as SplitParser - self.parser = SplitParser() - else: - from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser as CombinedParser - self.parser = CombinedParser() - + self.parser = Parser() self.prev_desire = np.zeros(self.constants.DESIRE_LEN, dtype=np.float32) - self.full_frames: dict = {} - self._blob_cache: dict = {} - nv12_info = get_nv12_info(cam_w, cam_h) - self.frame_buf_params = dict.fromkeys(self._vision_input_names, nv12_info) - yuv_size = self.frame_buf_params[self._road_key][3] - frame_tensor = Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize() - big_frame_tensor = Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize() - self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=frame_tensor, big_frame=big_frame_tensor) + if self.warp is not None: + self.full_frames = {k: Tensor(np.zeros(nv12_info[3], dtype=np.uint8), device=self.WARP_DEV).contiguous().realize() for k in self._vision_input_names} + self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[self._road_key], big_frame=self.full_frames[self._wide_key]) def warmup(self) -> None: - dummy_frames = {k: np.zeros(self.frame_buf_params[k][3], dtype=np.uint8) for k in self._vision_input_names} + dummy_size = self.frame_copy_size if self.is_run_model else self.frame_buf_params[self._road_key][3] + dummy_frames = {k: np.zeros(dummy_size, dtype=np.uint8) for k in self._vision_input_names} transforms = {k: np.eye(3, dtype=np.float32) for k in [self._road_key, self._wide_key] if k} - - dummy_inputs = {} - for k, v in self.numpy_inputs.items(): - if k not in ['tfm', 'big_tfm', 'prev_feat']: - dummy_inputs[k] = np.zeros(v.shape, dtype=v.dtype) - - self.run(dummy_frames, transforms, dummy_inputs, prepare_only=False) - - for v in self.numpy_inputs.values(): - v[:] = 0 + dummy_inputs = {k: np.zeros(v.shape, dtype=v.dtype) for k, v in self.numpy_inputs.items() if k not in ['tfm', 'big_tfm', 'prev_feat']} + self.run(dummy_frames, transforms, dummy_inputs) + if self.is_run_model: + self.input_queues, self.numpy_inputs, self.frame_buffers = make_stock_input_queues( + self.input_shapes, self.frame_skip, device=self.DEV, frame_copy_size=self.frame_copy_size) + self.frame_views = self.frame_buffers + self.npy = self.numpy_inputs + else: + for v in self.numpy_inputs.values(): + v[:] = 0 + self.full_frames.clear() + self._blob_cache.clear() self.prev_desire[:] = 0 - self.full_frames.clear() - self._blob_cache.clear() - @property def mlsim(self) -> bool: @@ -208,34 +214,40 @@ class ModelState(ModelStateBase): return self._desire_key def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], - inputs: dict[str, np.ndarray], prepare_only: bool, + inputs: dict[str, np.ndarray], after_enqueue: Callable[[], None] | None = None) -> dict[str, np.ndarray] | None: - for key in bufs.keys(): - ptr = np.frombuffer(bufs[key].data, dtype=np.uint8).ctypes.data - yuv_size = self.frame_buf_params[key][3] - cache_key = (key, ptr) - if cache_key not in self._blob_cache: - self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.WARP_DEV) - self.full_frames[key] = self._blob_cache[cache_key] + if self.is_run_model: + for key, buf in bufs.items(): + data = buf.data if hasattr(buf, 'data') else buf + np.copyto(self.frame_buffers[key], np.frombuffer(data, dtype=np.uint8, count=self.frame_copy_size)) + else: + for key, buf in bufs.items(): + ptr = np.frombuffer(buf.data, dtype=np.uint8).ctypes.data + cache_key = (key, ptr) + if cache_key not in self._blob_cache: + self._blob_cache[cache_key] = Tensor.from_blob(ptr, (self.frame_buf_params[key][3],), dtype='uint8', device=self.WARP_DEV) + self.full_frames[key] = self._blob_cache[cache_key] desire_key = self.desire_key inputs[desire_key][0] = 0 self.numpy_inputs[desire_key][:] = np.where(inputs[desire_key] - self.prev_desire > .99, inputs[desire_key], 0) self.prev_desire[:] = inputs[desire_key] + for key in ('traffic_convention', 'lateral_control_params', 'action_t'): if key in self.numpy_inputs and key in inputs: self.numpy_inputs[key][:] = inputs[key] - road_key = self._road_key - wide_key = self._wide_key - self.numpy_inputs['tfm'][:, :] = transforms[road_key].reshape(3, 3) - self.numpy_inputs['big_tfm'][:, :] = transforms[wide_key].reshape(3, 3) + self.numpy_inputs['tfm'][:, :] = transforms[self._road_key].reshape(3, 3) + self.numpy_inputs['big_tfm'][:, :] = transforms[self._wide_key].reshape(3, 3) + + if self.run_model is not None: + outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS}) + raw_outputs = outs + else: + assert self.warp is not None and self.run_policy is not None + warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[self._road_key], big_frame=self.full_frames[self._wide_key]) + raw_outputs = self.run_policy(**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped) - if prepare_only: - self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key]) - return None - warped = self.warp(**{k: self.input_queues[k] for k in WARP_INPUTS}, frame=self.full_frames[road_key], big_frame=self.full_frames[wide_key]) - raw_outputs = self.run_policy(**{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, warped=warped) if after_enqueue is not None: after_enqueue() @@ -245,7 +257,7 @@ class ModelState(ModelStateBase): raise RuntimeError("model output not finite") sliced = {k: model_output[np.newaxis, v] for k, v in self.vision_output_slices.items()} outputs = self.parser.parse_outputs(sliced) - if 'prev_feat' in self.numpy_inputs: + if 'prev_feat' in self.numpy_inputs and 'hidden_state' in self.vision_output_slices: self.numpy_inputs['prev_feat'][:] = model_output[self.vision_output_slices['hidden_state']] else: vision_output = raw_outputs[0].numpy().flatten() @@ -360,7 +372,11 @@ def main(demo=False): loader.start() loader.join(BIG_MODEL_TIMEOUT) model = big_model + if model is None: + params.put_bool("ChestnutModelError", True) params.put_bool("ChestnutActive", model is not None) + if model is not None: + params.remove("ChestnutModelError") small_model = ModelState(cam_w=vipc_client_main.width, cam_h=vipc_client_main.height, chestnut=False) if model is None or CHESTNUT else None if model is None: @@ -474,9 +490,6 @@ def main(demo=False): run_count = run_count + 1 frame_drop_ratio = frames_dropped / (1 + frames_dropped) - prepare_only = vipc_dropped_frames > 0 - if prepare_only: - cloudlog.error(f"skipping model eval. Dropped {vipc_dropped_frames} frames") bufs = {name: buf_extra if 'big' in name else buf_main for name in model.vision_input_names} transforms = {name: model_transform_extra if 'big' in name else model_transform_main for name in model.vision_input_names} @@ -501,11 +514,12 @@ def main(demo=False): try: send_chestnut = (chestnut_state is not None and run_count % round(model.constants.MODEL_FREQ / SERVICE_LIST['chestnutState'].frequency) == 0) - model_output = model.run(bufs, transforms, inputs, prepare_only, chestnut_state.send if send_chestnut else None) + model_output = model.run(bufs, transforms, inputs, chestnut_state.send if send_chestnut else None) except Exception: if not params.get_bool("ChestnutActive"): raise cloudlog.exception("chestnut failed, falling back to small") + params.put_bool("ChestnutModelError", True) params.put_bool("ChestnutActive", False) assert small_model is not None model = small_model diff --git a/openpilot/sunnypilot/modeld_v2/parse_model_outputs.py b/openpilot/sunnypilot/modeld_v2/parse_model_outputs.py index 7a3adcc1fa..d6cfdbbc88 100644 --- a/openpilot/sunnypilot/modeld_v2/parse_model_outputs.py +++ b/openpilot/sunnypilot/modeld_v2/parse_model_outputs.py @@ -115,22 +115,41 @@ class Parser: outs[name + '_stds'] = pred_std_final.reshape(final_shape) def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - # supercombo (4955 / 102) and newer variants (e.g. 990 / 144). - self.parse_mdn('plan', outs, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) - self.parse_mdn('lane_lines', outs, out_shape=(ModelConstants.NUM_LANE_LINES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH)) - self.parse_mdn('road_edges', outs, out_shape=(ModelConstants.NUM_ROAD_EDGES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH)) - self.parse_mdn('pose', outs, out_shape=(ModelConstants.POSE_WIDTH,)) - self.parse_mdn('road_transform', outs, out_shape=(ModelConstants.POSE_WIDTH,)) + if 'plan' in outs: + self.parse_mdn('plan', outs, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) + if 'planplus' in outs: + self.parse_mdn('planplus', outs, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) + if 'lane_lines' in outs: + self.parse_mdn('lane_lines', outs, out_shape=(ModelConstants.NUM_LANE_LINES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH)) + if 'road_edges' in outs: + self.parse_mdn('road_edges', outs, out_shape=(ModelConstants.NUM_ROAD_EDGES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH)) + if 'pose' in outs: + self.parse_mdn('pose', outs, out_shape=(ModelConstants.POSE_WIDTH,)) + if 'road_transform' in outs: + self.parse_mdn('road_transform', outs, out_shape=(ModelConstants.POSE_WIDTH,)) if 'sim_pose' in outs: self.parse_mdn('sim_pose', outs, out_shape=(ModelConstants.POSE_WIDTH,)) - self.parse_mdn('wide_from_device_euler', outs, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,)) - self.parse_mdn('lead', outs, out_shape=(ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)) + if 'wide_from_device_euler' in outs: + self.parse_mdn('wide_from_device_euler', outs, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,)) + if 'lead' in outs: + self.parse_mdn('lead', outs, out_shape=(ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)) if 'lat_planner_solution' in outs: self.parse_mdn('lat_planner_solution', outs, out_shape=(ModelConstants.IDX_N, ModelConstants.LAT_PLANNER_SOLUTION_WIDTH)) if 'desired_curvature' in outs: self.parse_mdn('desired_curvature', outs, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,)) + if 'action' in outs: + self.parse_mdn('action', outs, out_shape=(ModelConstants.ACTION_WIDTH,)) for k in ['lead_prob', 'lane_lines_prob', 'meta']: - self.parse_binary_crossentropy(k, outs) - self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,)) - self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN, ModelConstants.DESIRE_PRED_WIDTH)) + if k in outs: + self.parse_binary_crossentropy(k, outs) + if 'desire_state' in outs: + self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,)) + if 'desire_pred' in outs: + self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN, ModelConstants.DESIRE_PRED_WIDTH)) return outs + + def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + return self.parse_outputs(outs) + + def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + return self.parse_outputs(outs) diff --git a/openpilot/sunnypilot/modeld_v2/parse_model_outputs_split.py b/openpilot/sunnypilot/modeld_v2/parse_model_outputs_split.py deleted file mode 100644 index 3db47aee42..0000000000 --- a/openpilot/sunnypilot/modeld_v2/parse_model_outputs_split.py +++ /dev/null @@ -1,159 +0,0 @@ -import numpy as np -from openpilot.sunnypilot.models.split_model_constants import SplitModelConstants - - -def safe_exp(x, out=None): - # -11 is around 10**14, more causes float16 overflow - return np.exp(np.clip(x, -np.inf, 11), out=out) - - -def sigmoid(x): - return 1. / (1. + safe_exp(-x)) - - -def softmax(x, axis=-1): - x -= np.max(x, axis=axis, keepdims=True) - if x.dtype == np.float32 or x.dtype == np.float64: - safe_exp(x, out=x) - else: - x = safe_exp(x) - x /= np.sum(x, axis=axis, keepdims=True) - return x - - -class Parser: - def __init__(self, ignore_missing=False): - self.ignore_missing = ignore_missing - - def check_missing(self, outs, name): - if name not in outs and not self.ignore_missing: - raise ValueError(f"Missing output {name}") - return name not in outs - - def parse_categorical_crossentropy(self, name, outs, out_shape=None): - if self.check_missing(outs, name): - return - raw = outs[name] - if out_shape is not None: - raw = raw.reshape((raw.shape[0],) + out_shape) - outs[name] = softmax(raw, axis=-1) - - def parse_binary_crossentropy(self, name, outs): - if self.check_missing(outs, name): - return - raw = outs[name] - outs[name] = sigmoid(raw) - - def parse_mdn(self, name, outs, in_N=0, out_N=1, out_shape=None): - if self.check_missing(outs, name): - return - raw = outs[name] - raw = raw.reshape((raw.shape[0], max(in_N, 1), -1)) - - n_values = (raw.shape[2] - out_N)//2 - pred_mu = raw[:,:,:n_values] - pred_std = safe_exp(raw[:,:,n_values: 2*n_values]) - - if in_N > 1: - weights = np.zeros((raw.shape[0], in_N, out_N), dtype=raw.dtype) - for i in range(out_N): - weights[:,:,i - out_N] = softmax(raw[:,:,i - out_N], axis=-1) - - if out_N == 1: - for fidx in range(weights.shape[0]): - idxs = np.argsort(weights[fidx][:,0])[::-1] - weights[fidx] = weights[fidx][idxs] - pred_mu[fidx] = pred_mu[fidx][idxs] - pred_std[fidx] = pred_std[fidx][idxs] - assert out_shape is not None - full_shape = tuple([raw.shape[0], in_N] + list(out_shape)) - outs[name + '_weights'] = weights - outs[name + '_hypotheses'] = pred_mu.reshape(full_shape) - outs[name + '_stds_hypotheses'] = pred_std.reshape(full_shape) - - pred_mu_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) - pred_std_final = np.zeros((raw.shape[0], out_N, n_values), dtype=raw.dtype) - for fidx in range(weights.shape[0]): - for hidx in range(out_N): - idxs = np.argsort(weights[fidx,:,hidx])[::-1] - pred_mu_final[fidx, hidx] = pred_mu[fidx, idxs[0]] - pred_std_final[fidx, hidx] = pred_std[fidx, idxs[0]] - else: - pred_mu_final = pred_mu - pred_std_final = pred_std - - if out_N > 1: - assert out_shape is not None - final_shape = tuple([raw.shape[0], out_N] + list(out_shape)) - else: - assert out_shape is not None - final_shape = tuple([raw.shape[0],] + list(out_shape)) - outs[name] = pred_mu_final.reshape(final_shape) - outs[name + '_stds'] = pred_std_final.reshape(final_shape) - - def is_mhp(self, outs, name, shape): - if self.check_missing(outs, name): - return False - if outs[name].shape[1] == 2 * shape: - return False - return True - - def parse_dynamic_outputs(self, outs: dict[str, np.ndarray]) -> None: - if 'lead' in outs: - lead_mhp = self.is_mhp(outs, 'lead', - SplitModelConstants.LEAD_MHP_SELECTION * SplitModelConstants.LEAD_TRAJ_LEN * SplitModelConstants.LEAD_WIDTH) - lead_in_N, lead_out_N = (SplitModelConstants.LEAD_MHP_N, SplitModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0) - lead_out_shape = (SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) if lead_mhp else \ - (SplitModelConstants.LEAD_MHP_SELECTION, SplitModelConstants.LEAD_TRAJ_LEN, SplitModelConstants.LEAD_WIDTH) - self.parse_mdn('lead', outs, in_N=lead_in_N, out_N=lead_out_N, out_shape=lead_out_shape) - if 'plan' in outs: - plan_mhp = self.is_mhp(outs, 'plan', SplitModelConstants.IDX_N * SplitModelConstants.PLAN_WIDTH) - plan_in_N, plan_out_N = (SplitModelConstants.PLAN_MHP_N, SplitModelConstants.PLAN_MHP_SELECTION) if plan_mhp else (0, 0) - self.parse_mdn('plan', outs, in_N=plan_in_N, out_N=plan_out_N, - out_shape=(SplitModelConstants.IDX_N, SplitModelConstants.PLAN_WIDTH)) - if 'planplus' in outs: - self.parse_mdn('planplus', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.IDX_N, SplitModelConstants.PLAN_WIDTH)) - - def split_outputs(self, outs: dict[str, np.ndarray]) -> None: - if 'desired_curvature' in outs: - self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.DESIRED_CURV_WIDTH,)) - if 'desire_pred' in outs: - self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(SplitModelConstants.DESIRE_PRED_LEN,SplitModelConstants.DESIRE_PRED_WIDTH)) - if 'desire_state' in outs: - self.parse_categorical_crossentropy('desire_state', outs, out_shape=(SplitModelConstants.DESIRE_PRED_WIDTH,)) - if 'lane_lines' in outs: - self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, - out_shape=(SplitModelConstants.NUM_LANE_LINES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH)) - if 'lane_lines_prob' in outs: - self.parse_binary_crossentropy('lane_lines_prob', outs) - if 'lead_prob' in outs: - self.parse_binary_crossentropy('lead_prob', outs) - if 'lat_planner_solution' in outs: - self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.IDX_N,SplitModelConstants.LAT_PLANNER_SOLUTION_WIDTH)) - if 'meta' in outs: - self.parse_binary_crossentropy('meta', outs) - if 'road_edges' in outs: - self.parse_mdn('road_edges', outs, in_N=0, out_N=0, - out_shape=(SplitModelConstants.NUM_ROAD_EDGES,SplitModelConstants.IDX_N,SplitModelConstants.LANE_LINES_WIDTH)) - if 'sim_pose' in outs: - self.parse_mdn('sim_pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) - if 'action' in outs: - self.parse_mdn('action', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.ACTION_WIDTH,)) - - def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) - self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.WIDE_FROM_DEVICE_WIDTH,)) - self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(SplitModelConstants.POSE_WIDTH,)) - self.parse_dynamic_outputs(outs) - self.split_outputs(outs) - return outs - - def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - self.parse_dynamic_outputs(outs) - self.split_outputs(outs) - return outs - - def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - outs = self.parse_vision_outputs(outs) - outs = self.parse_policy_outputs(outs) - return outs diff --git a/openpilot/sunnypilot/modeld_v2/tests/helpers.py b/openpilot/sunnypilot/modeld_v2/tests/helpers.py index 82e159a305..8f148f6789 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/helpers.py +++ b/openpilot/sunnypilot/modeld_v2/tests/helpers.py @@ -117,7 +117,7 @@ ARCHETYPES = { is_20hz=True, expected_model_type='split', expected_constants_class=SplitModelConstants, - expected_parser_module='parse_model_outputs_split', + expected_parser_module='parse_model_outputs', expected_desire_key='desire', ), 'vision_multi_policy': Archetype( @@ -130,7 +130,7 @@ ARCHETYPES = { is_20hz=True, expected_model_type='multi_policy', expected_constants_class=SplitModelConstants, - expected_parser_module='parse_model_outputs_split', + expected_parser_module='parse_model_outputs', expected_desire_key='desire', ), 'tri_policy': Archetype( @@ -144,7 +144,7 @@ ARCHETYPES = { is_20hz=True, expected_model_type='multi_policy', expected_constants_class=SplitModelConstants, - expected_parser_module='parse_model_outputs_split', + expected_parser_module='parse_model_outputs', expected_desire_key='desire', ), 'supercombo_non20hz': Archetype( diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py b/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py index e3a686a73f..88cee8c5b4 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_combined_pkl_loader.py @@ -103,6 +103,23 @@ class TestStockEquivalence(OpenpilotTestCase): assert state.vision_output_slices == arch.metadata_structure['vision']['output_slices'] assert state.policy_output_slices == arch.metadata_structure['policy']['output_slices'] + def test_unified_run_model(self, tmp_path, monkeypatch, patch_modeld): + from openpilot.common.hardware import hw + from openpilot.selfdrive.modeld.helpers import dump_oob + shapes = {'img': (1, 12, 128, 256), 'big_img': (1, 12, 128, 256), 'features_buffer': (1, 24, 32, 512), + 'desire_pulse': (1, 25, 8), 'traffic_convention': (1, 2), 'action_t': (1, 2)} + pkl_data = {'metadata': {'model': {'input_shapes': shapes, 'output_slices': {}}}, + 'run_model': {(CAM_W, CAM_H): tests_helpers._noop_jit}} + with open(tmp_path / 'driving_test_tinygrad.pkl', 'wb') as f: + dump_oob(pkl_data, f) + bundle = DummyBundle(models=[DummyModel('supercombo', 'driving_test_tinygrad.pkl')]) + patch_modeld(bundle) + monkeypatch.setattr(hw.Paths, 'model_root', staticmethod(lambda: str(tmp_path))) + state = ModelState(cam_w=CAM_W, cam_h=CAM_H) + assert state.is_run_model and state.run_model is not None + assert state.run_policy is None and state.warp is None + assert 'img' in state.frame_views and 'big_img' in state.frame_views + ARCHETYPE_NAMES = list(ARCHETYPES.keys()) diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_parse_model_outputs.py b/openpilot/sunnypilot/modeld_v2/tests/test_parse_model_outputs.py new file mode 100644 index 0000000000..821043124a --- /dev/null +++ b/openpilot/sunnypilot/modeld_v2/tests/test_parse_model_outputs.py @@ -0,0 +1,81 @@ +import numpy as np +from openpilot.common.test import OpenpilotTestCase +from openpilot.sunnypilot.modeld_v2.constants import ModelConstants +from openpilot.sunnypilot.modeld_v2.parse_model_outputs import Parser, _infer_mhp, sigmoid, softmax + + +class TestParseModelOutputs(OpenpilotTestCase): + def test_infer_mhp_lead(self): + in_hypotheses, out_selections = _infer_mhp(102, 24) + assert in_hypotheses == 2 + assert out_selections == 3 + + def test_infer_mhp_plan(self): + in_hypotheses, out_selections = _infer_mhp(4955, 495) + assert in_hypotheses == 5 + assert out_selections == 1 + + def test_infer_mhp_non_mdn(self): + in_hypotheses, out_selections = _infer_mhp(48, 24) + assert in_hypotheses == 1 + assert out_selections == 0 + + def test_check_missing_raises(self): + parser = Parser(ignore_missing=False) + with self.assertRaises(ValueError): + parser.check_missing({}, "missing_key") + + def test_check_missing_ignored(self): + parser = Parser(ignore_missing=True) + assert parser.check_missing({}, "missing_key") is True + + def test_binary_crossentropy(self): + parser = Parser() + raw_logits = np.array([[-10.0, 0.0, 10.0]], dtype=np.float32) + outs = {"meta": raw_logits.copy()} + parser.parse_binary_crossentropy("meta", outs) + expected_probabilities = sigmoid(raw_logits) + np.testing.assert_allclose(outs["meta"], expected_probabilities, rtol=1e-5, atol=1e-6) + + def test_categorical_crossentropy(self): + parser = Parser() + raw_logits = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) + outs = {"desire_state": raw_logits.copy()} + parser.parse_categorical_crossentropy("desire_state", outs) + expected_probabilities = softmax(raw_logits) + np.testing.assert_allclose(outs["desire_state"], expected_probabilities, rtol=1e-5, atol=1e-6) + + def test_parse_vision_outputs(self): + parser = Parser() + pose_raw = np.zeros((1, ModelConstants.POSE_WIDTH * 2), dtype=np.float32) + road_transform_raw = np.zeros((1, ModelConstants.POSE_WIDTH * 2), dtype=np.float32) + lead_raw = np.zeros((1, 102), dtype=np.float32) + meta_raw = np.zeros((1, 55), dtype=np.float32) + vision_outputs = {"pose": pose_raw, "road_transform": road_transform_raw, "lead": lead_raw, "meta": meta_raw} + parsed = parser.parse_vision_outputs(vision_outputs) + assert "pose" in parsed + assert "road_transform" in parsed + assert "lead" in parsed + assert "meta" in parsed + assert parsed["pose"].shape == (1, ModelConstants.POSE_WIDTH) + assert parsed["lead"].shape == (1, ModelConstants.LEAD_MHP_SELECTION, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) + + def test_parse_policy_outputs(self): + parser = Parser() + plan_raw = np.zeros((1, 4955), dtype=np.float32) + desire_state_raw = np.zeros((1, ModelConstants.DESIRE_PRED_WIDTH), dtype=np.float32) + action_raw = np.zeros((1, ModelConstants.ACTION_WIDTH * 2), dtype=np.float32) + policy_outputs = {"plan": plan_raw, "desire_state": desire_state_raw, "action": action_raw} + parsed = parser.parse_policy_outputs(policy_outputs) + assert parsed["plan"].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + assert parsed["action"].shape == (1, ModelConstants.ACTION_WIDTH) + assert parsed["desire_state"].shape == (1, ModelConstants.DESIRE_PRED_WIDTH) + + def test_parse_outputs_combined(self): + parser = Parser() + outputs = {"plan": np.zeros((1, 4955), dtype=np.float32), "pose": np.zeros((1, ModelConstants.POSE_WIDTH * 2), + dtype=np.float32), "meta": np.zeros((1, 55), dtype=np.float32)} + parsed = parser.parse_outputs(outputs) + assert parsed["plan"].shape == (1, ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH) + assert parsed["pose"].shape == (1, ModelConstants.POSE_WIDTH) + assert parsed["meta"].shape == (1, 55) diff --git a/openpilot/sunnypilot/models/constants.py b/openpilot/sunnypilot/models/constants.py deleted file mode 100644 index 18abc6c960..0000000000 --- a/openpilot/sunnypilot/models/constants.py +++ /dev/null @@ -1,121 +0,0 @@ -import numpy as np - -def index_function(idx, max_val=192, max_idx=32): - return max_val * ((idx/max_idx)**2) - - -class ModelConstants: - # time and distance indices - IDX_N = 33 - T_IDXS = [index_function(idx, max_val=10.0) for idx in range(IDX_N)] - X_IDXS = [index_function(idx, max_val=192.0) for idx in range(IDX_N)] - LEAD_T_IDXS = [0., 2., 4., 6., 8., 10.] - LEAD_T_OFFSETS = [0., 2., 4.] - META_T_IDXS = [2., 4., 6., 8., 10.] - - # model inputs constants - MODEL_FREQ = 20 - FEATURE_LEN = 512 - HISTORY_BUFFER_LEN = 99 - DESIRE_LEN = 8 - TRAFFIC_CONVENTION_LEN = 2 - NAV_FEATURE_LEN = 256 - NAV_INSTRUCTION_LEN = 150 - LAT_PLANNER_STATE_LEN = 4 - LATERAL_CONTROL_PARAMS_LEN = 2 - PREV_DESIRED_CURV_LEN = 1 - - # model outputs constants - FCW_THRESHOLDS_5MS2 = np.array([.05, .05, .15, .15, .15], dtype=np.float32) - FCW_THRESHOLDS_3MS2 = np.array([.7, .7], dtype=np.float32) - FCW_5MS2_PROBS_WIDTH = 5 - FCW_3MS2_PROBS_WIDTH = 2 - - DISENGAGE_WIDTH = 5 - POSE_WIDTH = 6 - WIDE_FROM_DEVICE_WIDTH = 3 - SIM_POSE_WIDTH = 6 - LEAD_WIDTH = 4 - LANE_LINES_WIDTH = 2 - ROAD_EDGES_WIDTH = 2 - PLAN_WIDTH = 15 - DESIRE_PRED_WIDTH = 8 - LAT_PLANNER_SOLUTION_WIDTH = 4 - DESIRED_CURV_WIDTH = 1 - - NUM_LANE_LINES = 4 - NUM_ROAD_EDGES = 2 - - LEAD_TRAJ_LEN = 6 - DESIRE_PRED_LEN = 4 - - PLAN_MHP_N = 5 - LEAD_MHP_N = 2 - PLAN_MHP_SELECTION = 1 - LEAD_MHP_SELECTION = 3 - - FCW_THRESHOLD_5MS2_HIGH = 0.15 - FCW_THRESHOLD_5MS2_LOW = 0.05 - FCW_THRESHOLD_3MS2 = 0.7 - - CONFIDENCE_BUFFER_LEN = 5 - RYG_GREEN = 0.01165 - RYG_YELLOW = 0.06157 - - POLY_PATH_DEGREE = 4 - - -# model outputs slices -class Plan: - POSITION = slice(0, 3) - VELOCITY = slice(3, 6) - ACCELERATION = slice(6, 9) - T_FROM_CURRENT_EULER = slice(9, 12) - ORIENTATION_RATE = slice(12, 15) - - -class Meta: - ENGAGED = slice(0, 1) - # next 2, 4, 6, 8, 10 seconds - GAS_DISENGAGE = slice(1, 31, 6) - BRAKE_DISENGAGE = slice(2, 31, 6) - STEER_OVERRIDE = slice(3, 31, 6) - HARD_BRAKE_3 = slice(4, 31, 6) - HARD_BRAKE_4 = slice(5, 31, 6) - HARD_BRAKE_5 = slice(6, 31, 6) - # next 0, 2, 4, 6, 8, 10 seconds - GAS_PRESS = slice(31, 55, 4) - BRAKE_PRESS = slice(32, 55, 4) - LEFT_BLINKER = slice(33, 55, 4) - RIGHT_BLINKER = slice(34, 55, 4) - - -class MetaTombRaider: - ENGAGED = slice(0, 1) - # next 2, 4, 6, 8, 10 seconds - GAS_DISENGAGE = slice(1, 41, 8) - BRAKE_DISENGAGE = slice(2, 41, 8) - STEER_OVERRIDE = slice(3, 41, 8) - HARD_BRAKE_3 = slice(4, 41, 8) - HARD_BRAKE_4 = slice(5, 41, 8) - HARD_BRAKE_5 = slice(6, 41, 8) - GAS_PRESS = slice(7, 41, 8) - BRAKE_PRESS = slice(8, 41, 8) - # next 0, 2, 4, 6, 8, 10 seconds - LEFT_BLINKER = slice(41, 53, 2) - RIGHT_BLINKER = slice(42, 53, 2) - - -class MetaSimPose: - ENGAGED = slice(0, 1) - # next 2, 4, 6, 8, 10 seconds - GAS_DISENGAGE = slice(1, 36, 7) - BRAKE_DISENGAGE = slice(2, 36, 7) - STEER_OVERRIDE = slice(3, 36, 7) - HARD_BRAKE_3 = slice(4, 36, 7) - HARD_BRAKE_4 = slice(5, 36, 7) - HARD_BRAKE_5 = slice(6, 36, 7) - GAS_PRESS = slice(7, 36, 7) - # next 0, 2, 4, 6, 8, 10 seconds - LEFT_BLINKER = slice(36, 48, 2) - RIGHT_BLINKER = slice(37, 48, 2) diff --git a/openpilot/sunnypilot/models/fetcher.py b/openpilot/sunnypilot/models/fetcher.py index 810b9b7e49..16fa378c4e 100644 --- a/openpilot/sunnypilot/models/fetcher.py +++ b/openpilot/sunnypilot/models/fetcher.py @@ -139,7 +139,7 @@ class ModelCache: class ModelFetcher: """Handles fetching and caching of model data from remote source""" MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v22.json" - MODEL_URL_CHESTNUT = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_chestnut_v24.json" + MODEL_URL_CHESTNUT = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_chestnut_v25.json" MODEL_SOURCES = { "qcom": (MODEL_URL, ""), diff --git a/openpilot/sunnypilot/models/helpers.py b/openpilot/sunnypilot/models/helpers.py index ac69096597..a1a392c5ee 100644 --- a/openpilot/sunnypilot/models/helpers.py +++ b/openpilot/sunnypilot/models/helpers.py @@ -7,14 +7,11 @@ See the LICENSE.md file in the root directory for more details. import hashlib import os -import pickle -from pathlib import Path import numpy as np from openpilot.cereal import custom from openpilot.common.params import Params from openpilot.common.swaglog import cloudlog -from openpilot.sunnypilot.models.constants import Meta, MetaSimPose, MetaTombRaider from openpilot.common.hardware.hw import Paths from openpilot.selfdrive.modeld.helpers import chestnut_present @@ -22,7 +19,6 @@ from openpilot.selfdrive.modeld.helpers import chestnut_present REQUIRED_JSON_VERSION = 19 CUSTOM_MODEL_PATH = Paths.model_root() -METADATA_PATH = Path(__file__).parent / '../models/supercombo_metadata.pkl' ModelManager = custom.ModelManagerSP ACTIVE_BUNDLE_KEYS = { @@ -201,33 +197,6 @@ def _get_model(): return None -def load_metadata(): - metadata_path = METADATA_PATH - - with open(metadata_path, 'rb') as f: - return pickle.load(f) - - -def prepare_inputs(model_metadata: dict) -> dict[str, np.ndarray]: - return { - key: np.zeros(shape, dtype=np.float32).flatten() - for key, shape in model_metadata['input_shapes'].items() - if 'img' not in key - } - - -def load_meta_constants(model_metadata: dict): - """ Loads the appropriate meta model class based on key shapes""" - if 'sim_pose' in model_metadata['input_shapes']: - return MetaSimPose - - meta_slice = model_metadata['output_slices']['meta'] - if (meta_slice.start, meta_slice.stop, meta_slice.step) == (5868, 5921, None): - return MetaTombRaider - - return Meta - - # The following method(s) are modeld helper methods def plan_x_idxs_helper(constants, plan, model_output) -> list[float]: # times at X_IDXS according to plan.