diff --git a/openpilot/sunnypilot/modeld_v2/compile_modeld.py b/openpilot/sunnypilot/modeld_v2/compile_modeld.py index 85ae57c078..59056549f4 100755 --- a/openpilot/sunnypilot/modeld_v2/compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/compile_modeld.py @@ -7,6 +7,7 @@ See the LICENSE.md file in the root directory for more details. """ import argparse +import math import os import tempfile import time @@ -66,14 +67,15 @@ def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tu if desire_key: shapes['desire'] = (input_shapes[desire_key][2],) - if is_supercombo and 'features_buffer' in input_shapes: - fb = input_shapes['features_buffer'] - shapes['prev_feat'] = (fb[0], fb[2]) - for key, shape in input_shapes.items(): if key not in (desire_key, 'features_buffer') and 'img' not in key: shapes[key] = tuple(shape) + if is_supercombo and 'features_buffer' in input_shapes: + fb = input_shapes['features_buffer'] + feat_dim = math.prod(fb[2:]) + shapes['prev_feat'] = (fb[0], feat_dim) + sizes = [int(np.prod(size)) for size in shapes.values()] return shapes, sizes @@ -117,8 +119,9 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D } if features_buffer: + feat_dim = math.prod(features_buffer[2:]) feat_q_len = frame_skip * features_buffer[1] if is_supercombo else frame_skip * (features_buffer[1] - 1) + 1 - queues['feat_q'] = Tensor(np.zeros((feat_q_len, features_buffer[0], features_buffer[2]), + queues['feat_q'] = Tensor(np.zeros((feat_q_len, features_buffer[0], feat_dim), dtype=np.float32), device=device).contiguous().realize() queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items() if key in ('tfm', 'big_tfm')}) @@ -183,14 +186,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).realize() - big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn).realize() + 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) 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).realize() + desire_buf = 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(): @@ -199,19 +202,22 @@ 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).realize() + feat_buf = shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn) + inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb) 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() + feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize() + inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb) 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]) inputs.update({road_key: img, wide_key: big_img}) if 'features_buffer' not in inputs: - inputs['features_buffer'] = sample_skip_fn(feat_q) + feat_buf = sample_skip_fn(feat_q) + inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb) policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32').realize() if 'features_buffer' not in inputs and features_slice is not None: diff --git a/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py b/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py index 96bfb42638..86974b14f1 100644 --- a/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py +++ b/openpilot/sunnypilot/modeld_v2/tests/test_compile_modeld.py @@ -195,3 +195,85 @@ class TestReadFileChunkedToDisk(OpenpilotTestCase): assert out.parent == Path(d) assert out.read_bytes() == payload + + +class Test4DFeaturesBuffer(OpenpilotTestCase): + def test_get_policy_npy_shapes_4d(self): + from openpilot.sunnypilot.modeld_v2.compile_modeld import get_policy_npy_shapes + input_shapes = { + 'desire_pulse': (1, 25, 8), + 'features_buffer': (1, 24, 32, 512), # compare 4d to 3d for regression + 'traffic_convention': (1, 2), + 'action_t': (1, 2) + } + shapes, sizes = get_policy_npy_shapes(input_shapes, is_supercombo=True) + assert shapes['prev_feat'] == (1, 16384) + assert sizes == [8, 2, 2, 16384] + + def test_get_policy_npy_shapes_3d(self): + from openpilot.sunnypilot.modeld_v2.compile_modeld import get_policy_npy_shapes + input_shapes = { + 'desire_pulse': (1, 25, 8), + 'features_buffer': (1, 24, 512), + 'traffic_convention': (1, 2), + 'action_t': (1, 2) + } + shapes, sizes = get_policy_npy_shapes(input_shapes, is_supercombo=True) + assert shapes['prev_feat'] == (1, 512) + assert sizes == [8, 2, 2, 512] + + +class TestStockCompileModeldEquivalence(OpenpilotTestCase): + def test_get_policy_npy_shapes_matches_stock(self): + from openpilot.selfdrive.modeld.compile_modeld import get_policy_npy_shapes as stock_get_policy_npy_shapes + from openpilot.sunnypilot.modeld_v2.compile_modeld import get_policy_npy_shapes as sunny_get_policy_npy_shapes + + stock_input_shapes = { + 'desire_pulse': (1, 25, 8), + 'features_buffer': (1, 24, 512), # see below comment + 'traffic_convention': (1, 2), + 'action_t': (1, 2), + } + + stock_shapes, stock_sizes = stock_get_policy_npy_shapes(stock_input_shapes) + sunny_shapes, sunny_sizes = sunny_get_policy_npy_shapes(stock_input_shapes, is_supercombo=True) + + assert sunny_shapes == stock_shapes + assert sunny_sizes == stock_sizes + assert sunny_shapes['prev_feat'] == (1, 512) + + def test_make_input_queues_full_stock_equivalence(self): + from openpilot.selfdrive.modeld.compile_modeld import make_input_queues as stock_make_input_queues + from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues as sunny_make_supercombo_input_queues + input_shapes = { + 'img': (1, 12, 128, 256), + 'desire_pulse': (1, 25, 8), + 'features_buffer': (1, 24, 512), # when https://github.com/commaai/openpilot/pull/38681 merges, update to 1,24,32,512 + 'traffic_convention': (1, 2), + 'action_t': (1, 2), + } + frame_skip = 4 + + stock_queues, stock_npy = stock_make_input_queues(input_shapes, frame_skip, device='NPY') + sunny_queues, sunny_npy = sunny_make_supercombo_input_queues(input_shapes, frame_skip, device='NPY') + assert set(sunny_queues.keys()) == set(stock_queues.keys()) + for key in stock_queues: + assert sunny_queues[key].shape == stock_queues[key].shape, \ + f"Queue shape mismatch for {key}: sunny {sunny_queues[key].shape} != stock {stock_queues[key].shape}" + assert set(sunny_npy.keys()) == set(stock_npy.keys()) + for key in stock_npy: + assert sunny_npy[key].shape == stock_npy[key].shape, \ + f"Numpy array shape mismatch for {key}: sunny {sunny_npy[key].shape} != stock {stock_npy[key].shape}" + + def test_make_warp_queues_stock_equivalence(self): + from openpilot.selfdrive.modeld.compile_modeld import make_warp_input_queues as stock_make_warp_queues + from openpilot.sunnypilot.modeld_v2.compile_modeld import make_warp_queues as sunny_make_warp_queues + stock_vision_shapes = {'img': (1, 12, 128, 256)} # for now? + stock_queues, stock_npy = stock_make_warp_queues(stock_vision_shapes, frame_skip=4, device='NPY') + sunny_queues, sunny_npy = sunny_make_warp_queues(device='NPY') + + assert set(sunny_npy.keys()) == set(stock_npy.keys()) == {'tfm', 'big_tfm'} + for key in sunny_npy: + assert sunny_npy[key].shape == stock_npy[key].shape == (3, 3) + + diff --git a/openpilot/sunnypilot/models/fetcher.py b/openpilot/sunnypilot/models/fetcher.py index b64e27b1ad..5c23b8ef33 100644 --- a/openpilot/sunnypilot/models/fetcher.py +++ b/openpilot/sunnypilot/models/fetcher.py @@ -141,7 +141,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_v20.json" - MODEL_URL_USBGPU = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_usbgpu_v21.json" + MODEL_URL_USBGPU = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_usbgpu_v22.json" def __init__(self, params: Params): self.params = params