From ad5abd242ada0bcd89e279658dafb7d3663d358b Mon Sep 17 00:00:00 2001 From: discountchubbs Date: Sat, 6 Jun 2026 09:58:48 -0700 Subject: [PATCH] modeld_v2: refactor compile_modeld --- sunnypilot/modeld_v2/compile_modeld.py | 551 ++++++++----------------- 1 file changed, 167 insertions(+), 384 deletions(-) diff --git a/sunnypilot/modeld_v2/compile_modeld.py b/sunnypilot/modeld_v2/compile_modeld.py index 9c8cdb2c0a..b9a392631a 100755 --- a/sunnypilot/modeld_v2/compile_modeld.py +++ b/sunnypilot/modeld_v2/compile_modeld.py @@ -10,372 +10,199 @@ import argparse import os import pickle import time -from functools import partial from collections import defaultdict - +from functools import partial import numpy as np -from tinygrad.tensor import Tensor + +from tinygrad import dtypes from tinygrad.device import Device from tinygrad.engine.jit import TinyJit +from tinygrad.tensor import Tensor + +from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare, sample_desire, sample_skip, shift_and_sample -from openpilot.selfdrive.modeld.compile_modeld import ( - NV12Frame, make_frame_prepare, - shift_and_sample, sample_skip, sample_desire, -) MODEL_TYPES = ('vision_policy', 'supercombo', 'vision_multi_policy') -def _detect_desire_key(policy_input_shapes): - for k in policy_input_shapes: - if k.startswith('desire'): - return k - return None +def _detect_desire_key(shapes: dict) -> str | None: + return next((key for key in shapes if key.startswith('desire')), None) -def _detect_vision_keys(vision_input_shapes): - img_keys = sorted([k for k in vision_input_shapes if 'img' in k]) - road_key = next((k for k in img_keys if 'big' not in k), None) - wide_key = next((k for k in img_keys if 'big' in k), None) - if road_key is None or wide_key is None: - raise ValueError(f"Cannot determine road/wide image keys from {list(vision_input_shapes.keys())}") - return road_key, wide_key +def _detect_vision_keys(shapes: dict) -> tuple[str | None, str | None]: + img_keys = sorted(key for key in shapes if 'img' in key) + return ( + next((key for key in img_keys if 'big' not in key), None), + next((key for key in img_keys if 'big' in key), None) + ) -def make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device): - road_key, _ = _detect_vision_keys(vision_input_shapes) - img = vision_input_shapes[road_key] - n_frames = img[1] // 6 - img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3]) - - fb = policy_input_shapes['features_buffer'] - desire_key = _detect_desire_key(policy_input_shapes) - dp = policy_input_shapes[desire_key] - tc = policy_input_shapes.get('traffic_convention', (1, 2)) - - npy = { - 'desire': np.zeros(dp[2], dtype=np.float32), - 'traffic_convention': np.zeros(tc, dtype=np.float32), - 'tfm': np.zeros((3, 3), dtype=np.float32), - 'big_tfm': np.zeros((3, 3), dtype=np.float32), - } - - handled = {'features_buffer', desire_key, 'traffic_convention'} - for key, shape in policy_input_shapes.items(): - if key in handled: - continue - npy[key] = np.zeros(shape, dtype=np.float32) - - input_queues = { - 'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), - 'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), - 'feat_q': Tensor(np.zeros((frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]), dtype=np.float32), device=device).contiguous().realize(), - 'desire_q': Tensor(np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32), device=device).contiguous().realize(), - **{k: Tensor(v, device='NPY').realize() for k, v in npy.items()}, - } - return input_queues, npy +def derive_frame_skip(vision_input_shapes: dict, policy_input_shapes: dict) -> int: + features_buffer = policy_input_shapes.get('features_buffer') + return 1 if not features_buffer or features_buffer[1] >= 99 else 4 -def make_run_split_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, model_h, - vision_features_slice, frame_skip, desire_key, extra_policy_keys, - vision_road_key, vision_wide_key, prepare_only=False): - frame_prepare = make_frame_prepare(nv12, model_w, model_h) - sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) - sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) - - def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame, **extra): - npy_tensors = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT), - desire.to(Device.DEFAULT), traffic_convention.to(Device.DEFAULT)] - extra_device = {k: extra[k].to(Device.DEFAULT) for k in extra_policy_keys} - Tensor.realize(*npy_tensors, *extra_device.values()) - tfm, big_tfm, desire, traffic_convention = npy_tensors - - img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn) - big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn) - - if prepare_only: - return img, big_img - - vision_out = next(iter(vision_runner({vision_road_key: img, vision_wide_key: big_img}).values())).cast('float32') - - new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0) - feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn) - desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn) - - inputs = {'features_buffer': feat_buf, desire_key: desire_buf, 'traffic_convention': traffic_convention, **extra_device} - policy_out = next(iter(policy_runner(inputs).values())).cast('float32') - - return vision_out, policy_out - return run_policy - - -def compile_split_policy(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, - vision_runner, policy_runner, vision_metadata, policy_metadata): - print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") - - vision_features_slice = vision_metadata['output_slices']['hidden_state'] - vision_input_shapes = vision_metadata['input_shapes'] - policy_input_shapes = policy_metadata['input_shapes'] - desire_key = _detect_desire_key(policy_input_shapes) - extra_policy_keys = [k for k in policy_input_shapes if k not in ('features_buffer', desire_key, 'traffic_convention')] - vision_road_key, vision_wide_key = _detect_vision_keys(vision_input_shapes) - - _run = make_run_split_policy(vision_runner, policy_runner, nv12, model_w, model_h, - vision_features_slice, frame_skip, desire_key, extra_policy_keys, - vision_road_key, vision_wide_key, prepare_only) - run_policy_jit = TinyJit(_run, prune=True) - - SEED = 42 - - def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True): - input_queues, npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT) - np.random.seed(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): - frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() - big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() - for v in npy.values(): - v[:] = np.random.randn(*v.shape).astype(v.dtype) - Device.default.synchronize() - st = time.perf_counter() - outs = fn(**input_queues, frame=frame, big_frame=big_frame) - mt = time.perf_counter() - Device.default.synchronize() - et = time.perf_counter() - print(f" [{i+1}/{n_runs}] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") - - if i == 0: - val = [np.copy(v.numpy()) for v in outs] - buffers = [np.copy(v.numpy().copy()) for v in input_queues.values()] - - if test_val is not None: - match = all(np.array_equal(a, b) for a, b in zip(val, test_val, strict=True)) - assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed})" - if test_buffers is not None: - match = all(np.array_equal(a, b) for a, b in zip(buffers, test_buffers, strict=True)) - assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed})" - return fn, val, buffers - - print('capture + replay') - run_policy_jit, test_val, test_buffers = random_inputs_run_fn(run_policy_jit, SEED) - - print('pickle round trip') - run_policy_jit = pickle.loads(pickle.dumps(run_policy_jit)) - random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers, expect_match=True) - random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False) - return run_policy_jit - - -def derive_frame_skip(vision_input_shapes, policy_input_shapes): - fb = policy_input_shapes.get('features_buffer') - if fb is None: - return 1 - fb_history = fb[1] - if fb_history >= 99: - return 1 - return 4 - - -def make_supercombo_input_queues(input_shapes, frame_skip, device): - img_shape = input_shapes.get('img', input_shapes.get('input_imgs')) - if img_shape is None: - raise ValueError("No img input found in model shapes") +def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]: + road_key, _ = _detect_vision_keys(input_shapes) + if not road_key: + raise ValueError("Vision road key missing from input shapes.") + img_shape = input_shapes[road_key] n_frames = img_shape[1] // 6 img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img_shape[2], img_shape[3]) - numpy_keys = {} - queue_keys = {} + desire_key = _detect_desire_key(input_shapes) + if not desire_key: + raise ValueError("Desire key missing from input shapes.") + + desire_shape = input_shapes[desire_key] + features_buffer = input_shapes.get('features_buffer') + + npy_arrays = { + 'desire': np.zeros(desire_shape[2], dtype=np.float32), + 'tfm': np.zeros((3, 3), dtype=np.float32), + 'big_tfm': np.zeros((3, 3), dtype=np.float32) + } for key, shape in input_shapes.items(): - if 'img' in key: - continue - if len(shape) == 3 and shape[1] > 1: - if key.startswith('desire'): - numpy_keys[key] = np.zeros(shape[2], dtype=np.float32) - queue_keys[f'{key}_q'] = Tensor( - np.zeros((frame_skip * shape[1], shape[0], shape[2]), dtype=np.float32), - device=device).contiguous().realize() - elif key == 'features_buffer': - queue_keys['feat_q'] = Tensor( - np.zeros((frame_skip * (shape[1] - 1) + 1, shape[0], shape[2]), dtype=np.float32), - device=device).contiguous().realize() - else: - numpy_keys[key] = np.zeros(shape, dtype=np.float32) - elif len(shape) == 2: - numpy_keys[key] = np.zeros(shape, dtype=np.float32) + if key not in npy_arrays and 'img' not in key and key not in ('features_buffer', desire_key): + npy_arrays[key] = np.zeros(shape, dtype=np.float32) - if 'traffic_convention' not in numpy_keys: - tc_shape = input_shapes.get('traffic_convention', (1, 2)) - numpy_keys['traffic_convention'] = np.zeros(tc_shape, dtype=np.float32) - - numpy_keys['tfm'] = np.zeros((3, 3), dtype=np.float32) - numpy_keys['big_tfm'] = np.zeros((3, 3), dtype=np.float32) - - input_queues = { + queues = { 'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), 'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), - **queue_keys, - **{k: Tensor(v, device='NPY').realize() for k, v in numpy_keys.items()}, + 'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), + dtype=np.float32), device=device).contiguous().realize() } - return input_queues, numpy_keys + + if features_buffer: + queues['feat_q'] = Tensor(np.zeros((frame_skip * (features_buffer[1] - 1) + 1, features_buffer[0], features_buffer[2]), + dtype=np.float32), device=device).contiguous().realize() + + queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items()}) + return queues, npy_arrays -def make_run_supercombo(model_runner, nv12: NV12Frame, model_w, model_h, - features_slice, frame_skip, input_shapes, prepare_only=False): - frame_prepare = make_frame_prepare(nv12, model_w, model_h) +def make_split_input_queues(vision_input_shapes: dict, policy_input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]: + return generate_queues_and_npy({**vision_input_shapes, **policy_input_shapes}, frame_skip, device) + + +def make_supercombo_input_queues(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]: + return generate_queues_and_npy(input_shapes, frame_skip, device) + + +def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, model_size: tuple[int, int], + features_slice: slice, frame_skip: int, input_shapes: dict, prepare_only: bool): + frame_prepare = make_frame_prepare(nv12, *model_size) sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) desire_key = _detect_desire_key(input_shapes) - if desire_key is None: - raise ValueError(f"No desire* key found in input_shapes: {list(input_shapes.keys())}") - road_img_key, wide_img_key = _detect_vision_keys(input_shapes) - extra_policy_keys = [k for k in input_shapes - if k not in (desire_key, 'features_buffer', 'traffic_convention') - and 'img' not in k] + road_key, wide_key = _detect_vision_keys(input_shapes) - def run_supercombo(img_q, big_img_q, feat_q, desire_q, - frame, big_frame, **kwargs): - desire = kwargs.get(desire_key) + if not desire_key or not road_key or not wide_key: + raise ValueError("Missing required vision or desire keys in input shapes.") + + extra_keys = [key for key in input_shapes if key not in (desire_key, 'features_buffer', 'traffic_convention') and 'img' not in key] + + def runner(img_q, big_img_q, feat_q, frame, big_frame, tfm, big_tfm, **kwargs): + desire_q = kwargs['desire_q'] + desire = kwargs['desire'] traffic_convention = kwargs.get('traffic_convention') - tfm = kwargs['tfm'] - big_tfm = kwargs['big_tfm'] - tfm = tfm.to(Device.DEFAULT) - big_tfm = big_tfm.to(Device.DEFAULT) - desire = desire.to(Device.DEFAULT) - traffic_convention = traffic_convention.to(Device.DEFAULT) - Tensor.realize(tfm, big_tfm, desire, traffic_convention) + npys = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT), desire.to(Device.DEFAULT)] + if traffic_convention is not None: + npys.append(traffic_convention.to(Device.DEFAULT)) - img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn) - big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn) + extra_tensors = {key: kwargs[key].to(Device.DEFAULT) for key in extra_keys if key in kwargs} + Tensor.realize(*npys, *extra_tensors.values()) + + tfm_dev, big_tfm_dev, desire_dev = npys[:3] + traffic_conv_dev = npys[3] if traffic_convention is not None else None + + img = shift_and_sample(img_q, frame_prepare(frame, tfm_dev).unsqueeze(0), sample_skip_fn) + big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm_dev).unsqueeze(0), sample_skip_fn) if prepare_only: return img, big_img - desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn) - feat_buf = sample_skip_fn(feat_q) + desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn) + inputs = {desire_key: desire_buf, **extra_tensors} + if traffic_conv_dev is not None: + inputs['traffic_convention'] = traffic_conv_dev - inputs = {road_img_key: img, wide_img_key: big_img, - desire_key: desire_buf, 'features_buffer': feat_buf, - 'traffic_convention': traffic_convention} - for k in extra_policy_keys: - if k in kwargs: - inputs[k] = kwargs[k].to(Device.DEFAULT) + if vision_runner: + vision_out = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast('float32') + new_feat = vision_out[:, features_slice].reshape(1, -1).unsqueeze(0) + inputs['features_buffer'] = shift_and_sample(feat_q, new_feat, sample_skip_fn) + policy_outs = [next(iter(runner(inputs).values())).cast('float32') for runner in policy_runners] + return (vision_out, *policy_outs) if len(policy_outs) > 1 else (vision_out, policy_outs[0]) - model_out = next(iter(model_runner(inputs).values())).cast('float32') - - new_feat = model_out[:, features_slice].reshape(1, -1).unsqueeze(0) + inputs.update({road_key: img, wide_key: big_img, 'features_buffer': sample_skip_fn(feat_q)}) + policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32') + new_feat = policy_out[:, features_slice].reshape(1, -1).unsqueeze(0) shift_and_sample(feat_q, new_feat, sample_skip_fn) + return policy_out - return model_out - - return run_supercombo + return runner -def make_run_vision_multi_policy(vision_runner, policy_runners, nv12: NV12Frame, model_w, model_h, - vision_features_slice, frame_skip, desire_key, extra_policy_keys, - vision_road_key, vision_wide_key, prepare_only=False): - frame_prepare = make_frame_prepare(nv12, model_w, model_h) - sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) - sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) +def compile_and_warmup(nv12: NV12Frame, model_size: tuple[int, int], prepare_only: bool, frame_skip: int, vision_runner, policy_runners: list, metadata: dict): + print(f"Compiling combined JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") - def run_multi_policy(img_q, big_img_q, feat_q, desire_q, desire, - traffic_convention, tfm, big_tfm, frame, big_frame, **extra): - npy_tensors = [tfm.to(Device.DEFAULT), big_tfm.to(Device.DEFAULT), - desire.to(Device.DEFAULT), traffic_convention.to(Device.DEFAULT)] - extra_device = {k: extra[k].to(Device.DEFAULT) for k in extra_policy_keys} - Tensor.realize(*npy_tensors, *extra_device.values()) - tfm, big_tfm, desire, traffic_convention = npy_tensors + all_shapes = {key: value for meta in metadata.values() for key, value in meta['input_shapes'].items()} - img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn) - big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn) + feat_meta = metadata.get('vision') or metadata.get('model') or metadata.get('policy') + if not feat_meta: + raise ValueError("Could not find vision, model, or policy metadata.") - if prepare_only: - return img, big_img + features_slice = feat_meta['output_slices']['hidden_state'] - vision_out = next(iter(vision_runner({vision_road_key: img, vision_wide_key: big_img}).values())).cast('float32') + run_func = create_jit_runner(vision_runner, policy_runners, nv12, model_size, features_slice, frame_skip, all_shapes, prepare_only) + run_jit = TinyJit(run_func, prune=True) + queues, npy_arrays = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT) - new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0) - feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn) - desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn) - - inputs = {'features_buffer': feat_buf, desire_key: desire_buf, 'traffic_convention': traffic_convention, **extra_device} - - policy_outputs = [] - for runner in policy_runners: - policy_out = next(iter(runner(inputs).values())).cast('float32') - policy_outputs.append(policy_out) - - return (vision_out, *policy_outputs) - - return run_multi_policy - - -def _warmup_and_serialize(run_jit, input_queues, npy, nv12): for i in range(3): np.random.seed(42 + i) - frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() - big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() - for v in npy.values(): - v[:] = np.random.randn(*v.shape).astype(v.dtype) + frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8).realize() + big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8).realize() + for arr in npy_arrays.values(): + arr[:] = np.random.randn(*arr.shape).astype(arr.dtype) + Device.default.synchronize() - st = time.perf_counter() - run_jit(**input_queues, frame=frame, big_frame=big_frame) - mt = time.perf_counter() + start_time = time.perf_counter() + run_jit(**queues, frame=frame, big_frame=big_frame) + mid_time = time.perf_counter() Device.default.synchronize() - et = time.perf_counter() - print(f" [{i + 1}/3] enqueue {(mt - st) * 1e3:6.2f} ms -- total {(et - st) * 1e3:6.2f} ms") - return pickle.loads(pickle.dumps(run_jit)) + print(f" [{i + 1}/3] enqueue {(mid_time - start_time) * 1e3:6.2f} ms -- total {(time.perf_counter() - start_time) * 1e3:6.2f} ms") + + return pickle.loads(pickle.dumps(run_jit)) if not prepare_only else run_jit -def compile_supercombo(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, - model_runner, metadata): - print(f"Compiling combined supercombo JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") - - features_slice = metadata['output_slices']['hidden_state'] - input_shapes = metadata['input_shapes'] - - _run = make_run_supercombo(model_runner, nv12, model_w, model_h, - features_slice, frame_skip, input_shapes, prepare_only) - run_jit = TinyJit(_run, prune=True) - - input_queues, npy = make_supercombo_input_queues(input_shapes, frame_skip, Device.DEFAULT) - - run_jit = _warmup_and_serialize(run_jit, input_queues, npy, nv12) - return run_jit +def _parse_size(size_str: str) -> tuple[int, int]: + width, height = size_str.lower().split('x') + return int(width), int(height) -def compile_multi_policy(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, - vision_runner, policy_runners, vision_metadata, policy_metadata): - print(f"Compiling combined multi-policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") - - vision_features_slice = vision_metadata['output_slices']['hidden_state'] - vision_input_shapes = vision_metadata['input_shapes'] - policy_input_shapes = policy_metadata['input_shapes'] - desire_key = _detect_desire_key(policy_input_shapes) - extra_policy_keys = [k for k in policy_input_shapes if k not in ('features_buffer', desire_key, 'traffic_convention')] - vision_road_key, vision_wide_key = _detect_vision_keys(vision_input_shapes) - - _run = make_run_vision_multi_policy(vision_runner, policy_runners, nv12, model_w, model_h, - vision_features_slice, frame_skip, desire_key, extra_policy_keys, - vision_road_key, vision_wide_key, prepare_only) - run_jit = TinyJit(_run, prune=True) - - input_queues, npy = make_split_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT) - - run_jit = _warmup_and_serialize(run_jit, input_queues, npy, nv12) - return run_jit +def _compile_for_resolutions(camera_resolutions: list, model_size: tuple[int, int], frame_skip: int, + vision_runner, policy_runners: list, metadata: dict) -> dict: + return { + (cam_w, cam_h): { + name: compile_and_warmup(NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)), model_size, prepare_only, + frame_skip, vision_runner, policy_runners, metadata) + for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] + } + for cam_w, cam_h in camera_resolutions + } -def _parse_size(s): - w, h = s.lower().split('x') - return int(w), int(h) +def _load_policy_runners(args: argparse.Namespace) -> tuple[list, list]: + runners, keys = [], [] + for name, onnx_arg in [('policy', args.policy_onnx), ('off_policy', args.off_policy_onnx), ('on_policy', args.on_policy_onnx)]: + if onnx_arg: + runners.append(OnnxRunner(onnx_arg)) + keys.append(name) + return runners, keys if __name__ == "__main__": @@ -383,98 +210,54 @@ if __name__ == "__main__": from openpilot.system.camerad.cameras.nv12_info import get_nv12_info from openpilot.selfdrive.modeld.get_model_metadata import make_metadata_dict - p = argparse.ArgumentParser(description="Compile combined JIT pkl for sunnypilot modeld_v2") - p.add_argument('--model-type', choices=MODEL_TYPES, required=True) - p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH') - p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True) - p.add_argument('--frame-skip', type=int, default=None, help='frame skip value (auto-derived if not provided)') - p.add_argument('--output', required=True) + parser = argparse.ArgumentParser(description="Compile combined JIT pkl for sunnypilot modeld_v2") + parser.add_argument('--model-type', choices=MODEL_TYPES, required=True) + 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('--output', required=True) - p.add_argument('--vision-onnx', help='vision ONNX (for split models)') - p.add_argument('--policy-onnx', help='policy ONNX (for vision_policy)') - p.add_argument('--off-policy-onnx', help='off-policy ONNX (for vision_multi_policy)') - p.add_argument('--on-policy-onnx', help='on-policy ONNX (for vision_multi_policy)') - p.add_argument('--supercombo-onnx', help='supercombo ONNX (for supercombo)') + parser.add_argument('--vision-onnx', help='vision ONNX (for split models)') + parser.add_argument('--policy-onnx', help='policy ONNX (for vision_policy)') + parser.add_argument('--off-policy-onnx', help='off-policy ONNX (for vision_multi_policy)') + parser.add_argument('--on-policy-onnx', help='on-policy ONNX (for vision_multi_policy)') + parser.add_argument('--supercombo-onnx', help='supercombo ONNX (for supercombo)') - args = p.parse_args() - out = defaultdict(dict) + args = parser.parse_args() + output_data = defaultdict(dict) + + vision_runner = OnnxRunner(args.vision_onnx) if args.vision_onnx else None if args.model_type == 'vision_policy': - assert args.vision_onnx and args.policy_onnx - vision_runner = OnnxRunner(args.vision_onnx) - policy_runner = OnnxRunner(args.policy_onnx) - out['metadata']['vision'] = make_metadata_dict(args.vision_onnx) - out['metadata']['policy'] = make_metadata_dict(args.policy_onnx) - - frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip(out['metadata']['vision']['input_shapes'], - out['metadata']['policy']['input_shapes']) - - for cam_w, cam_h in args.camera_resolutions: - nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) - model_w, model_h = args.model_size - out[(cam_w, cam_h)] = { - name: compile_split_policy(nv12, model_w, model_h, prepare_only, frame_skip, - vision_runner, policy_runner, - out['metadata']['vision'], out['metadata']['policy']) - for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] - } - + 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': assert args.supercombo_onnx - model_runner = OnnxRunner(args.supercombo_onnx) - out['metadata']['model'] = make_metadata_dict(args.supercombo_onnx) - - frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip({}, out['metadata']['model']['input_shapes']) - - for cam_w, cam_h in args.camera_resolutions: - nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) - model_w, model_h = args.model_size - out[(cam_w, cam_h)] = { - name: compile_supercombo(nv12, model_w, model_h, prepare_only, frame_skip, - model_runner, out['metadata']['model']) - for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] - } - + policy_runners = [OnnxRunner(args.supercombo_onnx)] + output_data['metadata'] = {'model': make_metadata_dict(args.supercombo_onnx)} elif args.model_type == 'vision_multi_policy': - assert args.vision_onnx - vision_runner = OnnxRunner(args.vision_onnx) - out['metadata']['vision'] = make_metadata_dict(args.vision_onnx) + assert vision_runner + policy_runners, policy_names = _load_policy_runners(args) + output_data['metadata'] = {'vision': make_metadata_dict(args.vision_onnx)} + for name, runner_arg in zip(policy_names, [args.policy_onnx, args.off_policy_onnx, args.on_policy_onnx], strict=True): + if runner_arg: + output_data['metadata'][name] = make_metadata_dict(runner_arg) - policy_runners = [] - policy_onnxes = [] - if args.policy_onnx: - policy_onnxes.append(('policy', args.policy_onnx)) - if args.off_policy_onnx: - policy_onnxes.append(('off_policy', args.off_policy_onnx)) - if args.on_policy_onnx: - policy_onnxes.append(('on_policy', args.on_policy_onnx)) + first_policy_meta = output_data['metadata'].get('policy', output_data['metadata'].get('model', output_data['metadata'].get('off_policy', {}))) + vision_meta = output_data['metadata'].get('vision', {}) - for name, onnx_path in policy_onnxes: - runner = OnnxRunner(onnx_path) - policy_runners.append(runner) - out['metadata'][name] = make_metadata_dict(onnx_path) + derived_frame_skip = args.frame_skip or derive_frame_skip(vision_meta.get('input_shapes', {}), first_policy_meta.get('input_shapes', {})) + output_data.update(_compile_for_resolutions(args.camera_resolutions, args.model_size, derived_frame_skip, + vision_runner, policy_runners, output_data['metadata'])) - first_policy_key = policy_onnxes[0][0] - frame_skip = args.frame_skip if args.frame_skip is not None else derive_frame_skip(out['metadata']['vision']['input_shapes'], - out['metadata'][first_policy_key]['input_shapes']) + with open(args.output, "wb") as file: + pickle.dump(output_data, file) - for cam_w, cam_h in args.camera_resolutions: - nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) - model_w, model_h = args.model_size - out[(cam_w, cam_h)] = { - name: compile_multi_policy(nv12, model_w, model_h, prepare_only, frame_skip, - vision_runner, policy_runners, - out['metadata']['vision'], out['metadata'][first_policy_key]) - for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] - } - - with open(args.output, "wb") as f: - pickle.dump(out, f) pkl_size = os.path.getsize(args.output) print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)") from openpilot.common.file_chunker import chunk_file, get_chunk_targets chunk_targets = get_chunk_targets(args.output, pkl_size) chunk_file(args.output, chunk_targets) - num_chunks = len(chunk_targets) - 1 - print(f"Chunked into {num_chunks} file(s)") + print(f"Chunked into {len(chunk_targets) - 1} file(s)")