From aba815e162a4a25d83e5d4e142ec34997efdfda9 Mon Sep 17 00:00:00 2001 From: Bruce Wayne Date: Mon, 4 May 2026 19:14:08 -0700 Subject: [PATCH 1/5] Add op11 split driving models --- selfdrive/modeld/SConscript | 66 ++---- selfdrive/modeld/compile_warp.py | 198 ++++++++++++++++ selfdrive/modeld/constants.py | 1 + selfdrive/modeld/fill_model_msg.py | 5 +- selfdrive/modeld/modeld.py | 223 +++++++++++++----- .../modeld/models/big_driving_policy.onnx | 1 - .../modeld/models/big_driving_vision.onnx | 1 - .../modeld/models/driving_off_policy.onnx | 3 + .../modeld/models/driving_on_policy.onnx | 3 + selfdrive/modeld/models/driving_policy.onnx | 3 - selfdrive/modeld/models/driving_vision.onnx | 4 +- selfdrive/modeld/parse_model_outputs.py | 18 +- selfdrive/modeld/tinygrad_helpers.py | 12 + 13 files changed, 421 insertions(+), 117 deletions(-) create mode 100755 selfdrive/modeld/compile_warp.py delete mode 120000 selfdrive/modeld/models/big_driving_policy.onnx delete mode 120000 selfdrive/modeld/models/big_driving_vision.onnx create mode 100644 selfdrive/modeld/models/driving_off_policy.onnx create mode 100644 selfdrive/modeld/models/driving_on_policy.onnx delete mode 100644 selfdrive/modeld/models/driving_policy.onnx create mode 100644 selfdrive/modeld/tinygrad_helpers.py diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 40b30b8c1bc4d0..05045d09847b6d 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -1,23 +1,10 @@ import glob import json import os -from itertools import product from SCons.Script import Value from openpilot.common.file_chunker import chunk_file, get_chunk_paths -from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye -from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE -from openpilot.selfdrive.modeld.constants import ModelConstants -from openpilot.selfdrive.modeld.helpers import CompileConfig from tinygrad import Device -CAMERA_CONFIGS = [ - (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 - (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 -] -MODELD_CONFIGS = [CompileConfig(cam_w, cam_h, prepare_only, 'driving_') - for (cam_w, cam_h), prepare_only in product(CAMERA_CONFIGS, [True, False])] -DM_WARP_CONFIGS = [CompileConfig(cam_w, cam_h, True, 'dm_') for cam_w, cam_h in CAMERA_CONFIGS] - Import('env', 'arch') chunker_file = File("#common/file_chunker.py") lenv = env.Clone() @@ -29,6 +16,7 @@ tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + " def estimate_pickle_max_size(onnx_size): return 1.2 * onnx_size + 10 * 1024 * 1024 # 20% + 10MB is plenty +# THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689 # get fastest TG config available = set(Device.get_available_devices()) if 'CUDA' in available: @@ -36,10 +24,10 @@ if 'CUDA' in available: tg_flags = f'DEV={tg_backend}' elif 'QCOM' in available: tg_backend = 'QCOM' - tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1' + tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0' else: tg_backend = 'CPU' if arch == 'Darwin' else 'CPU:LLVM' - tg_flags = f'DEV={tg_backend}' + tg_flags = f'DEV={tg_backend} THREADS=0' def write_tg_compiled_flags(target, source, env): with open(str(target[0]), "w") as f: @@ -56,41 +44,23 @@ compiled_flags_node = lenv.Command( mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else '' # Get model metadata -for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: +for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']: fn = File(f"models/{model_name}").abspath script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)] cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [compiled_flags_node], cmd) -modeld_dir = Dir("#selfdrive/modeld").abspath -compile_modeld_script = [File(f"{modeld_dir}/compile_modeld.py")] -compile_dm_warp_script = [File(f"{modeld_dir}/compile_dm_warp.py")] -driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_policy']] -driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']] - -model_w, model_h = MEDMODEL_INPUT_SIZE -frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ -for cfg in MODELD_CONFIGS: - cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py ' - f'--model-size {model_w}x{model_h} ' - f'--nv12 {",".join(str(x) for x in cfg.nv12)} ' - f'--vision-onnx {File("models/driving_vision.onnx").abspath} ' - f'--policy-onnx {File("models/driving_policy.onnx").abspath} ' - f'--output {cfg.pkl_path} --frame-skip {frame_skip}' - + (' --prepare-only' if cfg.prepare_only else '')) - node = lenv.Command(cfg.pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + driving_metadata_deps + [chunker_file, compiled_flags_node], cmd) - onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps) - chunk_targets = get_chunk_paths(cfg.pkl_path, estimate_pickle_max_size(onnx_sizes_sum)) - def do_chunk(target, source, env, pkl=cfg.pkl_path, chunks=chunk_targets): - chunk_file(pkl, chunks) - lenv.Command(chunk_targets, node, do_chunk) - -dm_w, dm_h = DM_INPUT_SIZE -for cfg in DM_WARP_CONFIGS: - cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_dm_warp.py ' - f'--nv12 {",".join(str(x) for x in cfg.nv12)} --warp-to {dm_w}x{dm_h} ' - f'--output {cfg.pkl_path}') - lenv.Command(cfg.pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [compiled_flags_node], cmd) +image_flag = { + 'larch64': 'IMAGE=2', +}.get(arch, 'IMAGE=0') +script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)] +compile_warp_cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py ' +from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye +warp_targets = [] +for cam in [_ar_ox_fisheye, _os_fisheye]: + w, h = cam.width, cam.height + warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath] +lenv.Command(warp_targets, tinygrad_files + script_files + [compiled_flags_node], compile_warp_cmd) def tg_compile(flags, model_name): pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' @@ -101,7 +71,7 @@ def tg_compile(flags, model_name): compile_node = lenv.Command( pkl, [onnx_path] + tinygrad_files + [chunker_file, compiled_flags_node], - f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', + f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', ) def do_chunk(target, source, env): chunk_file(pkl, chunk_targets) @@ -111,4 +81,6 @@ def tg_compile(flags, model_name): do_chunk, ) -tg_compile(tg_flags, 'dmonitoring_model') +# Compile small models +for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']: + tg_compile(tg_flags, model_name) diff --git a/selfdrive/modeld/compile_warp.py b/selfdrive/modeld/compile_warp.py new file mode 100755 index 00000000000000..47511f2a2b6b23 --- /dev/null +++ b/selfdrive/modeld/compile_warp.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +import time +import pickle +import numpy as np +from pathlib import Path +from tinygrad.tensor import Tensor +from tinygrad.helpers import Context +from tinygrad.device import Device +from tinygrad.engine.jit import TinyJit + +from openpilot.system.camerad.cameras.nv12_info import get_nv12_info +from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE +from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye + +MODELS_DIR = Path(__file__).parent / 'models' + +CAMERA_CONFIGS = [ + (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 + (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 +] + +UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32) +UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX) + +IMG_BUFFER_SHAPE = (30, MEDMODEL_INPUT_SIZE[1] // 2, MEDMODEL_INPUT_SIZE[0] // 2) + + +def warp_pkl_path(w, h): + return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' + + +def dm_warp_pkl_path(w, h): + return MODELS_DIR / f'dm_warp_{w}x{h}_tinygrad.pkl' + + +def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad): + w_dst, h_dst = dst_shape + h_src, w_src = src_shape + + x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1) + y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1) + + # inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather) + src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2] + src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2] + src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2] + + src_x = src_x / src_w + src_y = src_y / src_w + + x_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int') + y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int') + idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped + + return src_flat[idx] + + +def frames_to_tensor(frames, model_w, model_h): + H = (frames.shape[0] * 2) // 3 + W = frames.shape[1] + in_img1 = Tensor.cat(frames[0:H:2, 0::2], + frames[1:H:2, 0::2], + frames[0:H:2, 1::2], + frames[1:H:2, 1::2], + frames[H:H+H//4].reshape((H//2, W//2)), + frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2)) + return in_img1 + + +def make_frame_prepare(cam_w, cam_h, model_w, model_h): + stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h) + uv_offset = stride * y_height + stride_pad = stride - cam_w + + def frame_prepare_tinygrad(input_frame, M_inv): + # UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling + M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]]) + # deinterleave NV12 UV plane (UVUV... -> separate U, V) + uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride) + with Context(SPLIT_REDUCEOP=0): + y = warp_perspective_tinygrad(input_frame[:cam_h*stride], + M_inv, (model_w, model_h), + (cam_h, cam_w), stride_pad).realize() + u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(), + M_inv_uv, (model_w//2, model_h//2), + (cam_h//2, cam_w//2), 0).realize() + v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(), + M_inv_uv, (model_w//2, model_h//2), + (cam_h//2, cam_w//2), 0).realize() + yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w)) + tensor = frames_to_tensor(yuv, model_w, model_h) + return tensor + return frame_prepare_tinygrad + + +def make_update_img_input(frame_prepare, model_w, model_h): + def update_img_input_tinygrad(frame_buffer, frame, M_inv): + M_inv = M_inv.to(Device.DEFAULT) + new_img = frame_prepare(frame, M_inv) + frame_buffer.assign(frame_buffer[6:].cat(new_img, dim=0).contiguous()) + return Tensor.cat(frame_buffer[:6], frame_buffer[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2) + return update_img_input_tinygrad + + +def make_update_both_imgs(frame_prepare, model_w, model_h): + update_img = make_update_img_input(frame_prepare, model_w, model_h) + + def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv, + calib_big_img_buffer, new_big_img, M_inv_big): + calib_img_pair = update_img(calib_img_buffer, new_img, M_inv) + calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big) + return calib_img_pair, calib_big_img_pair + return update_both_imgs_tinygrad + + +def make_warp_dm(cam_w, cam_h, dm_w, dm_h): + stride, y_height, _, _ = get_nv12_info(cam_w, cam_h) + stride_pad = stride - cam_w + + def warp_dm(input_frame, M_inv): + M_inv = M_inv.to(Device.DEFAULT) + result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w) + return result + return warp_dm + + +def compile_modeld_warp(cam_w, cam_h): + model_w, model_h = MEDMODEL_INPUT_SIZE + _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) + + print(f"Compiling modeld warp for {cam_w}x{cam_h}...") + + frame_prepare = make_frame_prepare(cam_w, cam_h, model_w, model_h) + update_both_imgs = make_update_both_imgs(frame_prepare, model_w, model_h) + update_img_jit = TinyJit(update_both_imgs, prune=True) + + full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() + big_full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() + for i in range(10): + img_inputs = [full_buffer, + Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), + Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] + big_img_inputs = [big_full_buffer, + Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), + Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] + inputs = img_inputs + big_img_inputs + Device.default.synchronize() + + st = time.perf_counter() + _ = update_img_jit(*inputs) + mt = time.perf_counter() + Device.default.synchronize() + et = time.perf_counter() + print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") + + pkl_path = warp_pkl_path(cam_w, cam_h) + with open(pkl_path, "wb") as f: + pickle.dump(update_img_jit, f) + print(f" Saved to {pkl_path}") + + jit = pickle.load(open(pkl_path, "rb")) + jit(*inputs) + + +def compile_dm_warp(cam_w, cam_h): + dm_w, dm_h = DM_INPUT_SIZE + _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) + + print(f"Compiling DM warp for {cam_w}x{cam_h}...") + + warp_dm = make_warp_dm(cam_w, cam_h, dm_w, dm_h) + warp_dm_jit = TinyJit(warp_dm, prune=True) + + for i in range(10): + inputs = [Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), + Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] + Device.default.synchronize() + st = time.perf_counter() + warp_dm_jit(*inputs) + mt = time.perf_counter() + Device.default.synchronize() + et = time.perf_counter() + print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") + + pkl_path = dm_warp_pkl_path(cam_w, cam_h) + with open(pkl_path, "wb") as f: + pickle.dump(warp_dm_jit, f) + print(f" Saved to {pkl_path}") + + +def run_and_save_pickle(): + for cam_w, cam_h in CAMERA_CONFIGS: + compile_modeld_warp(cam_w, cam_h) + compile_dm_warp(cam_w, cam_h) + + +if __name__ == "__main__": + run_and_save_pickle() diff --git a/selfdrive/modeld/constants.py b/selfdrive/modeld/constants.py index ff7e1d86006e83..0fb09262d0192e 100644 --- a/selfdrive/modeld/constants.py +++ b/selfdrive/modeld/constants.py @@ -38,6 +38,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/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py index 82c4c92b1d53c7..92a2dfa58d7f3a 100644 --- a/selfdrive/modeld/fill_model_msg.py +++ b/selfdrive/modeld/fill_model_msg.py @@ -125,7 +125,10 @@ def fill_model_msg(base_msg: capnp._DynamicStructBuilder, extended_msg: capnp._D # meta meta = modelV2.meta - meta.desireState = net_output_data['desire_state'][0].reshape(-1).tolist() + if 'desire_state' in net_output_data: + meta.desireState = net_output_data['desire_state'][0].reshape(-1).tolist() + else: + meta.desireState = [0.0] * ModelConstants.DESIRE_PRED_WIDTH meta.desirePrediction = net_output_data['desire_pred'][0].reshape(-1).tolist() meta.engagedProb = net_output_data['meta'][0,Meta.ENGAGED].item() meta.init('disengagePredictions') diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 73ed19ec943790..e543ccb1f728e5 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.helpers import MODELS_DIR, CompileConfig, set_tinygrad_backend_from_compiled_flags +from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags set_tinygrad_backend_from_compiled_flags() USBGPU = "USBGPU" in os.environ @@ -24,9 +24,8 @@ from openpilot.system.camerad.cameras.nv12_info import get_nv12_info from openpilot.common.transformations.model import get_warp_matrix from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper -from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value, get_curvature_from_plan +from openpilot.selfdrive.controls.lib.drive_helpers import MIN_SPEED, get_accel_from_plan, smooth_value, get_curvature_from_plan from openpilot.selfdrive.modeld.parse_model_outputs import Parser -from openpilot.selfdrive.modeld.compile_modeld import make_input_queues from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState from openpilot.common.file_chunker import read_file_chunked from openpilot.selfdrive.modeld.constants import ModelConstants, Plan @@ -35,33 +34,51 @@ PROCESS_NAME = "selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') +VISION_PKL_PATH = MODELS_DIR / 'driving_vision_tinygrad.pkl' VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl' -POLICY_METADATA_PATH = MODELS_DIR / 'driving_policy_metadata.pkl' +ON_POLICY_PKL_PATH = MODELS_DIR / 'driving_on_policy_tinygrad.pkl' +ON_POLICY_METADATA_PATH = MODELS_DIR / 'driving_on_policy_metadata.pkl' +OFF_POLICY_PKL_PATH = MODELS_DIR / 'driving_off_policy_tinygrad.pkl' +OFF_POLICY_METADATA_PATH = MODELS_DIR / 'driving_off_policy_metadata.pkl' LAT_SMOOTH_SECONDS = 0.0 LONG_SMOOTH_SECONDS = 0.3 MIN_LAT_CONTROL_SPEED = 0.3 +IMG_QUEUE_SHAPE = (6*(ModelConstants.MODEL_RUN_FREQ//ModelConstants.MODEL_CONTEXT_FREQ + 1), 128, 256) +assert IMG_QUEUE_SHAPE[0] == 30 def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action: plan = model_output['plan'][0] - desired_accel, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], - plan[:,Plan.ACCELERATION][:,0], - ModelConstants.T_IDXS, - action_t=long_action_t) - desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, LONG_SMOOTH_SECONDS) - - desired_curvature = get_curvature_from_plan(plan[:,Plan.T_FROM_CURRENT_EULER][:,2], - plan[:,Plan.ORIENTATION_RATE][:,2], - ModelConstants.T_IDXS, - v_ego, - lat_action_t) - if v_ego > MIN_LAT_CONTROL_SPEED: - desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, LAT_SMOOTH_SECONDS) + _, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], + plan[:,Plan.ACCELERATION][:,0], + ModelConstants.T_IDXS, + action_t=long_action_t) + + if 'action' in model_output: + desired_lat_accel, desired_accel = model_output['action'][0] + if v_ego > MIN_LAT_CONTROL_SPEED: + desired_curvature = desired_lat_accel / max(v_ego, MIN_SPEED) ** 2 + else: + desired_curvature = prev_action.desiredCurvature else: - desired_curvature = prev_action.desiredCurvature + desired_accel, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], + plan[:,Plan.ACCELERATION][:,0], + ModelConstants.T_IDXS, + action_t=long_action_t) + desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, LONG_SMOOTH_SECONDS) + + desired_curvature = get_curvature_from_plan(plan[:,Plan.T_FROM_CURRENT_EULER][:,2], + plan[:,Plan.ORIENTATION_RATE][:,2], + ModelConstants.T_IDXS, + v_ego, + lat_action_t) + if v_ego > MIN_LAT_CONTROL_SPEED: + desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, LAT_SMOOTH_SECONDS) + else: + desired_curvature = prev_action.desiredCurvature return log.ModelDataV2.Action(desiredCurvature=float(desired_curvature), desiredAcceleration=float(desired_accel), @@ -76,36 +93,114 @@ def __init__(self, vipc=None): if vipc is not None: self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof +class InputQueues: + def __init__ (self, model_fps, env_fps, n_frames_input): + assert env_fps % model_fps == 0 + assert env_fps >= model_fps + self.model_fps = model_fps + self.env_fps = env_fps + self.n_frames_input = n_frames_input + + self.dtypes = {} + self.shapes = {} + self.q = {} + + def update_dtypes_and_shapes(self, input_dtypes, input_shapes) -> None: + self.dtypes.update(input_dtypes) + if self.env_fps == self.model_fps: + self.shapes.update(input_shapes) + else: + for k in input_shapes: + shape = list(input_shapes[k]) + if 'img' in k: + n_channels = shape[1] // self.n_frames_input + shape[1] = (self.env_fps // self.model_fps + (self.n_frames_input - 1)) * n_channels + else: + shape[1] = (self.env_fps // self.model_fps) * shape[1] + self.shapes[k] = tuple(shape) + + def reset(self) -> None: + self.q = {k: np.zeros(self.shapes[k], dtype=self.dtypes[k]) for k in self.dtypes.keys()} + + def enqueue(self, inputs:dict[str, np.ndarray]) -> None: + for k in inputs.keys(): + if inputs[k].dtype != self.dtypes[k]: + raise ValueError(f'supplied input <{k}({inputs[k].dtype})> has wrong dtype, expected {self.dtypes[k]}') + input_shape = list(self.shapes[k]) + input_shape[1] = -1 + single_input = inputs[k].reshape(tuple(input_shape)) + sz = single_input.shape[1] + self.q[k][:,:-sz] = self.q[k][:,sz:] + self.q[k][:,-sz:] = single_input + + def get(self, *names) -> dict[str, np.ndarray]: + if self.env_fps == self.model_fps: + return {k: self.q[k] for k in names} + else: + out = {} + for k in names: + shape = self.shapes[k] + if 'img' in k: + n_channels = shape[1] // (self.env_fps // self.model_fps + (self.n_frames_input - 1)) + out[k] = np.concatenate([self.q[k][:, s:s+n_channels] for s in np.linspace(0, shape[1] - n_channels, self.n_frames_input, dtype=int)], axis=1) + elif 'pulse' in k: + # any pulse within interval counts + out[k] = self.q[k].reshape((shape[0], shape[1] * self.model_fps // self.env_fps, self.env_fps // self.model_fps, -1)).max(axis=2) + else: + idxs = np.arange(-1, -shape[1], -self.env_fps // self.model_fps)[::-1] + out[k] = self.q[k][:, idxs] + return out class ModelState: + inputs: dict[str, np.ndarray] + output: np.ndarray prev_desire: np.ndarray # for tracking the rising edge of the pulse - def __init__(self, cam_w: int, cam_h: int): + def __init__(self): with open(VISION_METADATA_PATH, 'rb') as f: vision_metadata = pickle.load(f) self.vision_input_shapes = vision_metadata['input_shapes'] self.vision_input_names = list(self.vision_input_shapes.keys()) self.vision_output_slices = vision_metadata['output_slices'] + vision_output_size = vision_metadata['output_shapes']['outputs'][1] + + with open(OFF_POLICY_METADATA_PATH, 'rb') as f: + off_policy_metadata = pickle.load(f) + self.off_policy_input_shapes = off_policy_metadata['input_shapes'] + self.off_policy_output_slices = off_policy_metadata['output_slices'] + off_policy_output_size = off_policy_metadata['output_shapes']['outputs'][1] - with open(POLICY_METADATA_PATH, 'rb') as f: + with open(ON_POLICY_METADATA_PATH, 'rb') as f: policy_metadata = pickle.load(f) self.policy_input_shapes = policy_metadata['input_shapes'] self.policy_output_slices = policy_metadata['output_slices'] + policy_output_size = policy_metadata['output_shapes']['outputs'][1] self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32) - self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ - self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip) + # policy inputs + self.numpy_inputs = {k: np.zeros(self.policy_input_shapes[k], dtype=np.float32) for k in self.policy_input_shapes} + self.full_input_queues = InputQueues(ModelConstants.MODEL_CONTEXT_FREQ, ModelConstants.MODEL_RUN_FREQ, ModelConstants.N_FRAMES) + for k in ['desire_pulse', 'features_buffer']: + self.full_input_queues.update_dtypes_and_shapes({k: self.numpy_inputs[k].dtype}, {k: self.numpy_inputs[k].shape}) + self.full_input_queues.reset() + + self.img_queues = {'img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize(), + 'big_img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize()} self.full_frames : dict[str, Tensor] = {} self._blob_cache : dict[int, Tensor] = {} + self.transforms_np = {k: np.zeros((3,3), dtype=np.float32) for k in self.img_queues} + self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()} + self.vision_output = np.zeros(vision_output_size, dtype=np.float32) + self.policy_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} + self.policy_output = np.zeros(policy_output_size, dtype=np.float32) + self.off_policy_output = np.zeros(off_policy_output_size, dtype=np.float32) self.parser = Parser() - self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')} - self.run_policy = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=False).pkl_path)) - self.warp_enqueue = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=True).pkl_path)) - self.warp_enqueue( - **self.input_queues, - frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(), - big_frame=Tensor.zeros(self.frame_buf_params['big_img'][3], dtype='uint8').contiguous().realize()) + self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {} + self.update_imgs = None + self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH))) + self.policy_run = pickle.loads(read_file_chunked(str(ON_POLICY_PKL_PATH))) + self.off_policy_run = pickle.loads(read_file_chunked(str(OFF_POLICY_PKL_PATH))) def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]: parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()} @@ -113,6 +208,18 @@ def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slic def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None: + # Model decides when action is completed, so desire input is just a pulse triggered on rising edge + inputs['desire_pulse'][0] = 0 + new_desire = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) + self.prev_desire[:] = inputs['desire_pulse'] + if self.update_imgs is None: + for key in bufs.keys(): + w, h = bufs[key].width, bufs[key].height + self.frame_buf_params[key] = get_nv12_info(w, h) + warp_path = MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' + with open(warp_path, "rb") as f: + self.update_imgs = pickle.load(f) + for key in bufs.keys(): ptr = bufs[key].data.ctypes.data yuv_size = self.frame_buf_params[key][3] @@ -121,31 +228,36 @@ def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], if cache_key not in self._blob_cache: self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8') self.full_frames[key] = self._blob_cache[cache_key] + for key in bufs.keys(): + self.transforms_np[key][:,:] = transforms[key][:,:] - # Model decides when action is completed, so desire input is just a pulse triggered on rising edge - inputs['desire_pulse'][0] = 0 - self.npy['desire'][:] = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) - self.prev_desire[:] = inputs['desire_pulse'] - self.npy['traffic_convention'][:] = inputs['traffic_convention'] - self.npy['tfm'][:,:] = transforms['img'][:,:] - self.npy['big_tfm'][:,:] = transforms['big_img'][:,:] + out = self.update_imgs(self.img_queues['img'], self.full_frames['img'], self.transforms['img'], + self.img_queues['big_img'], self.full_frames['big_img'], self.transforms['big_img']) + vision_inputs = {'img': out[0], 'big_img': out[1]} if prepare_only: - self.warp_enqueue(**self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img']) return None - vision_output, policy_output = self.run_policy( - **self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'] - ) + self.vision_output = self.vision_run(**vision_inputs).contiguous().realize().uop.base.buffer.numpy().flatten() + vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(self.vision_output, self.vision_output_slices)) - vision_output = vision_output.numpy().flatten() - policy_output = policy_output.numpy().flatten() - vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices)) - policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(policy_output, self.policy_output_slices)) - combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict} + self.full_input_queues.enqueue({'features_buffer': vision_outputs_dict['hidden_state'], 'desire_pulse': new_desire}) + for k in ['desire_pulse', 'features_buffer']: + self.numpy_inputs[k][:] = self.full_input_queues.get(k)[k] + self.numpy_inputs['traffic_convention'][:] = inputs['traffic_convention'] + if 'action_t' in self.numpy_inputs: + self.numpy_inputs['action_t'][:] = inputs['action_t'] + self.policy_output = self.policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy().flatten() + policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(self.policy_output, self.policy_output_slices)) + + self.off_policy_output = self.off_policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy() + off_policy_outputs_dict = self.parser.parse_off_policy_outputs(self.slice_outputs(self.off_policy_output, self.off_policy_output_slices)) + + combined_outputs_dict = {**vision_outputs_dict, **off_policy_outputs_dict, **policy_outputs_dict} if SEND_RAW_PRED: - combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), policy_output.copy()]) + combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy(), self.off_policy_output.copy()]) + return combined_outputs_dict @@ -157,6 +269,11 @@ def main(demo=False): # also need to move the aux USB interrupts for good timings config_realtime_process(7, 54) + st = time.monotonic() + cloudlog.warning("loading model") + model = ModelState() + cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") + # visionipc clients while True: available_streams = VisionIpcClient.available_streams("camerad", block=False) @@ -180,11 +297,6 @@ def main(demo=False): if use_extra_client: cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})") - st = time.monotonic() - cloudlog.warning("loading model") - model = ModelState(vipc_client_main.width, vipc_client_main.height) - cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") - # messaging pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry"]) sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"]) @@ -287,9 +399,14 @@ def main(demo=False): 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} + frame_delay = DT_MDL # compensate for time passed since the frame was captured: current_time - timestamp_eof is 50ms on average + action_delay = DT_MDL / 2 # middle of the interval between model output (current state) and next frame (expected state) + lat_action_t = lat_delay + frame_delay + action_delay + long_action_t = long_delay + frame_delay + action_delay inputs:dict[str, np.ndarray] = { 'desire_pulse': vec_desire, 'traffic_convention': traffic_convention, + 'action_t': np.array([lat_action_t, long_action_t], dtype=np.float32), } mt1 = time.perf_counter() @@ -302,9 +419,7 @@ def main(demo=False): drivingdata_send = messaging.new_message('drivingModelData') posenet_send = messaging.new_message('cameraOdometry') - frame_delay = DT_MDL # compensate for time passed since the frame was captured: current_time - timestamp_eof is 50ms on average - action_delay = DT_MDL / 2 # middle of the interval between model output (current state) and next frame (expected state) - action = get_action_from_model(model_output, prev_action, lat_delay + frame_delay + action_delay, long_delay + frame_delay + action_delay, v_ego) + action = get_action_from_model(model_output, prev_action, lat_action_t, long_action_t, v_ego) prev_action = action fill_model_msg(drivingdata_send, modelv2_send, model_output, action, publish_state, meta_main.frame_id, meta_extra.frame_id, frame_id, diff --git a/selfdrive/modeld/models/big_driving_policy.onnx b/selfdrive/modeld/models/big_driving_policy.onnx deleted file mode 120000 index e1b653a14a03d6..00000000000000 --- a/selfdrive/modeld/models/big_driving_policy.onnx +++ /dev/null @@ -1 +0,0 @@ -driving_policy.onnx \ No newline at end of file diff --git a/selfdrive/modeld/models/big_driving_vision.onnx b/selfdrive/modeld/models/big_driving_vision.onnx deleted file mode 120000 index 28ee71dd746e63..00000000000000 --- a/selfdrive/modeld/models/big_driving_vision.onnx +++ /dev/null @@ -1 +0,0 @@ -driving_vision.onnx \ No newline at end of file diff --git a/selfdrive/modeld/models/driving_off_policy.onnx b/selfdrive/modeld/models/driving_off_policy.onnx new file mode 100644 index 00000000000000..2975a571dcb392 --- /dev/null +++ b/selfdrive/modeld/models/driving_off_policy.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9a2bdfb988ab96fc60d991fd9e6b38221819b235539bb0d6b781fc029438f599 +size 13928996 diff --git a/selfdrive/modeld/models/driving_on_policy.onnx b/selfdrive/modeld/models/driving_on_policy.onnx new file mode 100644 index 00000000000000..239249a81c8f82 --- /dev/null +++ b/selfdrive/modeld/models/driving_on_policy.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0e66a9a3f5eb8da0b7693904ff8f1904b5fb43c9c0e265c170efad30c84629cb +size 12548985 diff --git a/selfdrive/modeld/models/driving_policy.onnx b/selfdrive/modeld/models/driving_policy.onnx deleted file mode 100644 index 611ae9fe85f837..00000000000000 --- a/selfdrive/modeld/models/driving_policy.onnx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:78477124cbf3ffe30fa951ebada8410b43c4242c6054584d656f1d329b067e15 -size 14060847 diff --git a/selfdrive/modeld/models/driving_vision.onnx b/selfdrive/modeld/models/driving_vision.onnx index 6c9fc4c84d3632..c01d6e233b8087 100644 --- a/selfdrive/modeld/models/driving_vision.onnx +++ b/selfdrive/modeld/models/driving_vision.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:ee29ee5bce84d1ce23e9ff381280de9b4e4d96d2934cd751740354884e112c66 -size 46877473 +oid sha256:5928713f355d75cf01ec4961bb0442c3712581dd3bdcf9aaca387eee77049f69 +size 23272727 diff --git a/selfdrive/modeld/parse_model_outputs.py b/selfdrive/modeld/parse_model_outputs.py index a0b45d2a981685..3211668b5aa5cc 100644 --- a/selfdrive/modeld/parse_model_outputs.py +++ b/selfdrive/modeld/parse_model_outputs.py @@ -96,11 +96,17 @@ def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndar self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,)) self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) + self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH)) + self.parse_binary_crossentropy('meta', outs) + return outs + + def parse_off_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + plan_mhp = self.is_mhp(outs, 'plan', ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH) + plan_in_N, plan_out_N = (ModelConstants.PLAN_MHP_N, ModelConstants.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=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH)) self.parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_ROAD_EDGES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH)) self.parse_binary_crossentropy('lane_lines_prob', outs) - self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH)) - self.parse_binary_crossentropy('meta', outs) self.parse_binary_crossentropy('lead_prob', outs) lead_mhp = self.is_mhp(outs, 'lead', ModelConstants.LEAD_MHP_SELECTION * ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH) lead_in_N, lead_out_N = (ModelConstants.LEAD_MHP_N, ModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0) @@ -110,15 +116,11 @@ def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndar return outs def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - plan_mhp = self.is_mhp(outs, 'plan', ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH) - plan_in_N, plan_out_N = (ModelConstants.PLAN_MHP_N, ModelConstants.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=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) - if 'planplus' in outs: - self.parse_mdn('planplus', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) - self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,)) + self.parse_mdn('action', outs, in_N=0, out_N=0, out_shape=(ModelConstants.ACTION_WIDTH,)) return outs def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: outs = self.parse_vision_outputs(outs) + outs = self.parse_off_policy_outputs(outs) outs = self.parse_policy_outputs(outs) return outs diff --git a/selfdrive/modeld/tinygrad_helpers.py b/selfdrive/modeld/tinygrad_helpers.py new file mode 100644 index 00000000000000..49a6ed6161855c --- /dev/null +++ b/selfdrive/modeld/tinygrad_helpers.py @@ -0,0 +1,12 @@ +import json +import os +from pathlib import Path + +MODELS_DIR = Path(__file__).parent / 'models' +COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json' + + +def set_tinygrad_backend_from_compiled_flags() -> None: + if os.path.isfile(COMPILED_FLAGS_PATH): + with open(COMPILED_FLAGS_PATH) as f: + os.environ['DEV'] = str(json.load(f)['DEV']) From a30a35a88f46bff7a41819b0cff9731a7894af9a Mon Sep 17 00:00:00 2001 From: Bruce Wayne Date: Mon, 4 May 2026 19:30:31 -0700 Subject: [PATCH 2/5] Revert "Add op11 split driving models" This reverts commit 223cefc7025d7da0ed6b2f934e155aa47b1da9bb. --- selfdrive/modeld/SConscript | 66 ++++-- selfdrive/modeld/compile_warp.py | 198 ---------------- selfdrive/modeld/constants.py | 1 - selfdrive/modeld/fill_model_msg.py | 5 +- selfdrive/modeld/modeld.py | 223 +++++------------- .../modeld/models/big_driving_policy.onnx | 1 + .../modeld/models/big_driving_vision.onnx | 1 + .../modeld/models/driving_off_policy.onnx | 3 - .../modeld/models/driving_on_policy.onnx | 3 - selfdrive/modeld/models/driving_policy.onnx | 3 + selfdrive/modeld/models/driving_vision.onnx | 4 +- selfdrive/modeld/parse_model_outputs.py | 18 +- selfdrive/modeld/tinygrad_helpers.py | 12 - 13 files changed, 117 insertions(+), 421 deletions(-) delete mode 100755 selfdrive/modeld/compile_warp.py create mode 120000 selfdrive/modeld/models/big_driving_policy.onnx create mode 120000 selfdrive/modeld/models/big_driving_vision.onnx delete mode 100644 selfdrive/modeld/models/driving_off_policy.onnx delete mode 100644 selfdrive/modeld/models/driving_on_policy.onnx create mode 100644 selfdrive/modeld/models/driving_policy.onnx delete mode 100644 selfdrive/modeld/tinygrad_helpers.py diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 05045d09847b6d..40b30b8c1bc4d0 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -1,10 +1,23 @@ import glob import json import os +from itertools import product from SCons.Script import Value from openpilot.common.file_chunker import chunk_file, get_chunk_paths +from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye +from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE +from openpilot.selfdrive.modeld.constants import ModelConstants +from openpilot.selfdrive.modeld.helpers import CompileConfig from tinygrad import Device +CAMERA_CONFIGS = [ + (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 + (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 +] +MODELD_CONFIGS = [CompileConfig(cam_w, cam_h, prepare_only, 'driving_') + for (cam_w, cam_h), prepare_only in product(CAMERA_CONFIGS, [True, False])] +DM_WARP_CONFIGS = [CompileConfig(cam_w, cam_h, True, 'dm_') for cam_w, cam_h in CAMERA_CONFIGS] + Import('env', 'arch') chunker_file = File("#common/file_chunker.py") lenv = env.Clone() @@ -16,7 +29,6 @@ tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + " def estimate_pickle_max_size(onnx_size): return 1.2 * onnx_size + 10 * 1024 * 1024 # 20% + 10MB is plenty -# THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689 # get fastest TG config available = set(Device.get_available_devices()) if 'CUDA' in available: @@ -24,10 +36,10 @@ if 'CUDA' in available: tg_flags = f'DEV={tg_backend}' elif 'QCOM' in available: tg_backend = 'QCOM' - tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0' + tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1' else: tg_backend = 'CPU' if arch == 'Darwin' else 'CPU:LLVM' - tg_flags = f'DEV={tg_backend} THREADS=0' + tg_flags = f'DEV={tg_backend}' def write_tg_compiled_flags(target, source, env): with open(str(target[0]), "w") as f: @@ -44,23 +56,41 @@ compiled_flags_node = lenv.Command( mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else '' # Get model metadata -for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']: +for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: fn = File(f"models/{model_name}").abspath script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)] cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [compiled_flags_node], cmd) -image_flag = { - 'larch64': 'IMAGE=2', -}.get(arch, 'IMAGE=0') -script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)] -compile_warp_cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py ' -from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye -warp_targets = [] -for cam in [_ar_ox_fisheye, _os_fisheye]: - w, h = cam.width, cam.height - warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath] -lenv.Command(warp_targets, tinygrad_files + script_files + [compiled_flags_node], compile_warp_cmd) +modeld_dir = Dir("#selfdrive/modeld").abspath +compile_modeld_script = [File(f"{modeld_dir}/compile_modeld.py")] +compile_dm_warp_script = [File(f"{modeld_dir}/compile_dm_warp.py")] +driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_policy']] +driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']] + +model_w, model_h = MEDMODEL_INPUT_SIZE +frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ +for cfg in MODELD_CONFIGS: + cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py ' + f'--model-size {model_w}x{model_h} ' + f'--nv12 {",".join(str(x) for x in cfg.nv12)} ' + f'--vision-onnx {File("models/driving_vision.onnx").abspath} ' + f'--policy-onnx {File("models/driving_policy.onnx").abspath} ' + f'--output {cfg.pkl_path} --frame-skip {frame_skip}' + + (' --prepare-only' if cfg.prepare_only else '')) + node = lenv.Command(cfg.pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + driving_metadata_deps + [chunker_file, compiled_flags_node], cmd) + onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps) + chunk_targets = get_chunk_paths(cfg.pkl_path, estimate_pickle_max_size(onnx_sizes_sum)) + def do_chunk(target, source, env, pkl=cfg.pkl_path, chunks=chunk_targets): + chunk_file(pkl, chunks) + lenv.Command(chunk_targets, node, do_chunk) + +dm_w, dm_h = DM_INPUT_SIZE +for cfg in DM_WARP_CONFIGS: + cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_dm_warp.py ' + f'--nv12 {",".join(str(x) for x in cfg.nv12)} --warp-to {dm_w}x{dm_h} ' + f'--output {cfg.pkl_path}') + lenv.Command(cfg.pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [compiled_flags_node], cmd) def tg_compile(flags, model_name): pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' @@ -71,7 +101,7 @@ def tg_compile(flags, model_name): compile_node = lenv.Command( pkl, [onnx_path] + tinygrad_files + [chunker_file, compiled_flags_node], - f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', + f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', ) def do_chunk(target, source, env): chunk_file(pkl, chunk_targets) @@ -81,6 +111,4 @@ def tg_compile(flags, model_name): do_chunk, ) -# Compile small models -for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']: - tg_compile(tg_flags, model_name) +tg_compile(tg_flags, 'dmonitoring_model') diff --git a/selfdrive/modeld/compile_warp.py b/selfdrive/modeld/compile_warp.py deleted file mode 100755 index 47511f2a2b6b23..00000000000000 --- a/selfdrive/modeld/compile_warp.py +++ /dev/null @@ -1,198 +0,0 @@ -#!/usr/bin/env python3 -import time -import pickle -import numpy as np -from pathlib import Path -from tinygrad.tensor import Tensor -from tinygrad.helpers import Context -from tinygrad.device import Device -from tinygrad.engine.jit import TinyJit - -from openpilot.system.camerad.cameras.nv12_info import get_nv12_info -from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE -from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye - -MODELS_DIR = Path(__file__).parent / 'models' - -CAMERA_CONFIGS = [ - (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 - (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 -] - -UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32) -UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX) - -IMG_BUFFER_SHAPE = (30, MEDMODEL_INPUT_SIZE[1] // 2, MEDMODEL_INPUT_SIZE[0] // 2) - - -def warp_pkl_path(w, h): - return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' - - -def dm_warp_pkl_path(w, h): - return MODELS_DIR / f'dm_warp_{w}x{h}_tinygrad.pkl' - - -def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad): - w_dst, h_dst = dst_shape - h_src, w_src = src_shape - - x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1) - y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1) - - # inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather) - src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2] - src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2] - src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2] - - src_x = src_x / src_w - src_y = src_y / src_w - - x_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int') - y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int') - idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped - - return src_flat[idx] - - -def frames_to_tensor(frames, model_w, model_h): - H = (frames.shape[0] * 2) // 3 - W = frames.shape[1] - in_img1 = Tensor.cat(frames[0:H:2, 0::2], - frames[1:H:2, 0::2], - frames[0:H:2, 1::2], - frames[1:H:2, 1::2], - frames[H:H+H//4].reshape((H//2, W//2)), - frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2)) - return in_img1 - - -def make_frame_prepare(cam_w, cam_h, model_w, model_h): - stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h) - uv_offset = stride * y_height - stride_pad = stride - cam_w - - def frame_prepare_tinygrad(input_frame, M_inv): - # UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling - M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]]) - # deinterleave NV12 UV plane (UVUV... -> separate U, V) - uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride) - with Context(SPLIT_REDUCEOP=0): - y = warp_perspective_tinygrad(input_frame[:cam_h*stride], - M_inv, (model_w, model_h), - (cam_h, cam_w), stride_pad).realize() - u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(), - M_inv_uv, (model_w//2, model_h//2), - (cam_h//2, cam_w//2), 0).realize() - v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(), - M_inv_uv, (model_w//2, model_h//2), - (cam_h//2, cam_w//2), 0).realize() - yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w)) - tensor = frames_to_tensor(yuv, model_w, model_h) - return tensor - return frame_prepare_tinygrad - - -def make_update_img_input(frame_prepare, model_w, model_h): - def update_img_input_tinygrad(frame_buffer, frame, M_inv): - M_inv = M_inv.to(Device.DEFAULT) - new_img = frame_prepare(frame, M_inv) - frame_buffer.assign(frame_buffer[6:].cat(new_img, dim=0).contiguous()) - return Tensor.cat(frame_buffer[:6], frame_buffer[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2) - return update_img_input_tinygrad - - -def make_update_both_imgs(frame_prepare, model_w, model_h): - update_img = make_update_img_input(frame_prepare, model_w, model_h) - - def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv, - calib_big_img_buffer, new_big_img, M_inv_big): - calib_img_pair = update_img(calib_img_buffer, new_img, M_inv) - calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big) - return calib_img_pair, calib_big_img_pair - return update_both_imgs_tinygrad - - -def make_warp_dm(cam_w, cam_h, dm_w, dm_h): - stride, y_height, _, _ = get_nv12_info(cam_w, cam_h) - stride_pad = stride - cam_w - - def warp_dm(input_frame, M_inv): - M_inv = M_inv.to(Device.DEFAULT) - result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w) - return result - return warp_dm - - -def compile_modeld_warp(cam_w, cam_h): - model_w, model_h = MEDMODEL_INPUT_SIZE - _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) - - print(f"Compiling modeld warp for {cam_w}x{cam_h}...") - - frame_prepare = make_frame_prepare(cam_w, cam_h, model_w, model_h) - update_both_imgs = make_update_both_imgs(frame_prepare, model_w, model_h) - update_img_jit = TinyJit(update_both_imgs, prune=True) - - full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() - big_full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() - for i in range(10): - img_inputs = [full_buffer, - Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), - Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - big_img_inputs = [big_full_buffer, - Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), - Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - inputs = img_inputs + big_img_inputs - Device.default.synchronize() - - st = time.perf_counter() - _ = update_img_jit(*inputs) - mt = time.perf_counter() - Device.default.synchronize() - et = time.perf_counter() - print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") - - pkl_path = warp_pkl_path(cam_w, cam_h) - with open(pkl_path, "wb") as f: - pickle.dump(update_img_jit, f) - print(f" Saved to {pkl_path}") - - jit = pickle.load(open(pkl_path, "rb")) - jit(*inputs) - - -def compile_dm_warp(cam_w, cam_h): - dm_w, dm_h = DM_INPUT_SIZE - _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) - - print(f"Compiling DM warp for {cam_w}x{cam_h}...") - - warp_dm = make_warp_dm(cam_w, cam_h, dm_w, dm_h) - warp_dm_jit = TinyJit(warp_dm, prune=True) - - for i in range(10): - inputs = [Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), - Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - Device.default.synchronize() - st = time.perf_counter() - warp_dm_jit(*inputs) - mt = time.perf_counter() - Device.default.synchronize() - et = time.perf_counter() - print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") - - pkl_path = dm_warp_pkl_path(cam_w, cam_h) - with open(pkl_path, "wb") as f: - pickle.dump(warp_dm_jit, f) - print(f" Saved to {pkl_path}") - - -def run_and_save_pickle(): - for cam_w, cam_h in CAMERA_CONFIGS: - compile_modeld_warp(cam_w, cam_h) - compile_dm_warp(cam_w, cam_h) - - -if __name__ == "__main__": - run_and_save_pickle() diff --git a/selfdrive/modeld/constants.py b/selfdrive/modeld/constants.py index 0fb09262d0192e..ff7e1d86006e83 100644 --- a/selfdrive/modeld/constants.py +++ b/selfdrive/modeld/constants.py @@ -38,7 +38,6 @@ 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/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py index 92a2dfa58d7f3a..82c4c92b1d53c7 100644 --- a/selfdrive/modeld/fill_model_msg.py +++ b/selfdrive/modeld/fill_model_msg.py @@ -125,10 +125,7 @@ def fill_model_msg(base_msg: capnp._DynamicStructBuilder, extended_msg: capnp._D # meta meta = modelV2.meta - if 'desire_state' in net_output_data: - meta.desireState = net_output_data['desire_state'][0].reshape(-1).tolist() - else: - meta.desireState = [0.0] * ModelConstants.DESIRE_PRED_WIDTH + meta.desireState = net_output_data['desire_state'][0].reshape(-1).tolist() meta.desirePrediction = net_output_data['desire_pred'][0].reshape(-1).tolist() meta.engagedProb = net_output_data['meta'][0,Meta.ENGAGED].item() meta.init('disengagePredictions') diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index e543ccb1f728e5..73ed19ec943790 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags +from openpilot.selfdrive.modeld.helpers import MODELS_DIR, CompileConfig, set_tinygrad_backend_from_compiled_flags set_tinygrad_backend_from_compiled_flags() USBGPU = "USBGPU" in os.environ @@ -24,8 +24,9 @@ from openpilot.system.camerad.cameras.nv12_info import get_nv12_info from openpilot.common.transformations.model import get_warp_matrix from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper -from openpilot.selfdrive.controls.lib.drive_helpers import MIN_SPEED, get_accel_from_plan, smooth_value, get_curvature_from_plan +from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value, get_curvature_from_plan from openpilot.selfdrive.modeld.parse_model_outputs import Parser +from openpilot.selfdrive.modeld.compile_modeld import make_input_queues from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState from openpilot.common.file_chunker import read_file_chunked from openpilot.selfdrive.modeld.constants import ModelConstants, Plan @@ -34,51 +35,33 @@ PROCESS_NAME = "selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') -VISION_PKL_PATH = MODELS_DIR / 'driving_vision_tinygrad.pkl' VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl' -ON_POLICY_PKL_PATH = MODELS_DIR / 'driving_on_policy_tinygrad.pkl' -ON_POLICY_METADATA_PATH = MODELS_DIR / 'driving_on_policy_metadata.pkl' -OFF_POLICY_PKL_PATH = MODELS_DIR / 'driving_off_policy_tinygrad.pkl' -OFF_POLICY_METADATA_PATH = MODELS_DIR / 'driving_off_policy_metadata.pkl' +POLICY_METADATA_PATH = MODELS_DIR / 'driving_policy_metadata.pkl' LAT_SMOOTH_SECONDS = 0.0 LONG_SMOOTH_SECONDS = 0.3 MIN_LAT_CONTROL_SPEED = 0.3 -IMG_QUEUE_SHAPE = (6*(ModelConstants.MODEL_RUN_FREQ//ModelConstants.MODEL_CONTEXT_FREQ + 1), 128, 256) -assert IMG_QUEUE_SHAPE[0] == 30 def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action: plan = model_output['plan'][0] - _, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], - plan[:,Plan.ACCELERATION][:,0], - ModelConstants.T_IDXS, - action_t=long_action_t) - - if 'action' in model_output: - desired_lat_accel, desired_accel = model_output['action'][0] - if v_ego > MIN_LAT_CONTROL_SPEED: - desired_curvature = desired_lat_accel / max(v_ego, MIN_SPEED) ** 2 - else: - desired_curvature = prev_action.desiredCurvature + desired_accel, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], + plan[:,Plan.ACCELERATION][:,0], + ModelConstants.T_IDXS, + action_t=long_action_t) + desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, LONG_SMOOTH_SECONDS) + + desired_curvature = get_curvature_from_plan(plan[:,Plan.T_FROM_CURRENT_EULER][:,2], + plan[:,Plan.ORIENTATION_RATE][:,2], + ModelConstants.T_IDXS, + v_ego, + lat_action_t) + if v_ego > MIN_LAT_CONTROL_SPEED: + desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, LAT_SMOOTH_SECONDS) else: - desired_accel, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], - plan[:,Plan.ACCELERATION][:,0], - ModelConstants.T_IDXS, - action_t=long_action_t) - desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, LONG_SMOOTH_SECONDS) - - desired_curvature = get_curvature_from_plan(plan[:,Plan.T_FROM_CURRENT_EULER][:,2], - plan[:,Plan.ORIENTATION_RATE][:,2], - ModelConstants.T_IDXS, - v_ego, - lat_action_t) - if v_ego > MIN_LAT_CONTROL_SPEED: - desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, LAT_SMOOTH_SECONDS) - else: - desired_curvature = prev_action.desiredCurvature + desired_curvature = prev_action.desiredCurvature return log.ModelDataV2.Action(desiredCurvature=float(desired_curvature), desiredAcceleration=float(desired_accel), @@ -93,114 +76,36 @@ def __init__(self, vipc=None): if vipc is not None: self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof -class InputQueues: - def __init__ (self, model_fps, env_fps, n_frames_input): - assert env_fps % model_fps == 0 - assert env_fps >= model_fps - self.model_fps = model_fps - self.env_fps = env_fps - self.n_frames_input = n_frames_input - - self.dtypes = {} - self.shapes = {} - self.q = {} - - def update_dtypes_and_shapes(self, input_dtypes, input_shapes) -> None: - self.dtypes.update(input_dtypes) - if self.env_fps == self.model_fps: - self.shapes.update(input_shapes) - else: - for k in input_shapes: - shape = list(input_shapes[k]) - if 'img' in k: - n_channels = shape[1] // self.n_frames_input - shape[1] = (self.env_fps // self.model_fps + (self.n_frames_input - 1)) * n_channels - else: - shape[1] = (self.env_fps // self.model_fps) * shape[1] - self.shapes[k] = tuple(shape) - - def reset(self) -> None: - self.q = {k: np.zeros(self.shapes[k], dtype=self.dtypes[k]) for k in self.dtypes.keys()} - - def enqueue(self, inputs:dict[str, np.ndarray]) -> None: - for k in inputs.keys(): - if inputs[k].dtype != self.dtypes[k]: - raise ValueError(f'supplied input <{k}({inputs[k].dtype})> has wrong dtype, expected {self.dtypes[k]}') - input_shape = list(self.shapes[k]) - input_shape[1] = -1 - single_input = inputs[k].reshape(tuple(input_shape)) - sz = single_input.shape[1] - self.q[k][:,:-sz] = self.q[k][:,sz:] - self.q[k][:,-sz:] = single_input - - def get(self, *names) -> dict[str, np.ndarray]: - if self.env_fps == self.model_fps: - return {k: self.q[k] for k in names} - else: - out = {} - for k in names: - shape = self.shapes[k] - if 'img' in k: - n_channels = shape[1] // (self.env_fps // self.model_fps + (self.n_frames_input - 1)) - out[k] = np.concatenate([self.q[k][:, s:s+n_channels] for s in np.linspace(0, shape[1] - n_channels, self.n_frames_input, dtype=int)], axis=1) - elif 'pulse' in k: - # any pulse within interval counts - out[k] = self.q[k].reshape((shape[0], shape[1] * self.model_fps // self.env_fps, self.env_fps // self.model_fps, -1)).max(axis=2) - else: - idxs = np.arange(-1, -shape[1], -self.env_fps // self.model_fps)[::-1] - out[k] = self.q[k][:, idxs] - return out class ModelState: - inputs: dict[str, np.ndarray] - output: np.ndarray prev_desire: np.ndarray # for tracking the rising edge of the pulse - def __init__(self): + def __init__(self, cam_w: int, cam_h: int): with open(VISION_METADATA_PATH, 'rb') as f: vision_metadata = pickle.load(f) self.vision_input_shapes = vision_metadata['input_shapes'] self.vision_input_names = list(self.vision_input_shapes.keys()) self.vision_output_slices = vision_metadata['output_slices'] - vision_output_size = vision_metadata['output_shapes']['outputs'][1] - - with open(OFF_POLICY_METADATA_PATH, 'rb') as f: - off_policy_metadata = pickle.load(f) - self.off_policy_input_shapes = off_policy_metadata['input_shapes'] - self.off_policy_output_slices = off_policy_metadata['output_slices'] - off_policy_output_size = off_policy_metadata['output_shapes']['outputs'][1] - with open(ON_POLICY_METADATA_PATH, 'rb') as f: + with open(POLICY_METADATA_PATH, 'rb') as f: policy_metadata = pickle.load(f) self.policy_input_shapes = policy_metadata['input_shapes'] self.policy_output_slices = policy_metadata['output_slices'] - policy_output_size = policy_metadata['output_shapes']['outputs'][1] self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32) - # policy inputs - self.numpy_inputs = {k: np.zeros(self.policy_input_shapes[k], dtype=np.float32) for k in self.policy_input_shapes} - self.full_input_queues = InputQueues(ModelConstants.MODEL_CONTEXT_FREQ, ModelConstants.MODEL_RUN_FREQ, ModelConstants.N_FRAMES) - for k in ['desire_pulse', 'features_buffer']: - self.full_input_queues.update_dtypes_and_shapes({k: self.numpy_inputs[k].dtype}, {k: self.numpy_inputs[k].shape}) - self.full_input_queues.reset() - - self.img_queues = {'img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize(), - 'big_img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize()} + self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ + self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip) self.full_frames : dict[str, Tensor] = {} self._blob_cache : dict[int, Tensor] = {} - self.transforms_np = {k: np.zeros((3,3), dtype=np.float32) for k in self.img_queues} - self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()} - self.vision_output = np.zeros(vision_output_size, dtype=np.float32) - self.policy_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} - self.policy_output = np.zeros(policy_output_size, dtype=np.float32) - self.off_policy_output = np.zeros(off_policy_output_size, dtype=np.float32) self.parser = Parser() - self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {} - self.update_imgs = None - self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH))) - self.policy_run = pickle.loads(read_file_chunked(str(ON_POLICY_PKL_PATH))) - self.off_policy_run = pickle.loads(read_file_chunked(str(OFF_POLICY_PKL_PATH))) + self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')} + self.run_policy = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=False).pkl_path)) + self.warp_enqueue = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=True).pkl_path)) + self.warp_enqueue( + **self.input_queues, + frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(), + big_frame=Tensor.zeros(self.frame_buf_params['big_img'][3], dtype='uint8').contiguous().realize()) def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]: parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()} @@ -208,18 +113,6 @@ def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slic def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None: - # Model decides when action is completed, so desire input is just a pulse triggered on rising edge - inputs['desire_pulse'][0] = 0 - new_desire = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) - self.prev_desire[:] = inputs['desire_pulse'] - if self.update_imgs is None: - for key in bufs.keys(): - w, h = bufs[key].width, bufs[key].height - self.frame_buf_params[key] = get_nv12_info(w, h) - warp_path = MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' - with open(warp_path, "rb") as f: - self.update_imgs = pickle.load(f) - for key in bufs.keys(): ptr = bufs[key].data.ctypes.data yuv_size = self.frame_buf_params[key][3] @@ -228,36 +121,31 @@ def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], if cache_key not in self._blob_cache: self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8') self.full_frames[key] = self._blob_cache[cache_key] - for key in bufs.keys(): - self.transforms_np[key][:,:] = transforms[key][:,:] - out = self.update_imgs(self.img_queues['img'], self.full_frames['img'], self.transforms['img'], - self.img_queues['big_img'], self.full_frames['big_img'], self.transforms['big_img']) - vision_inputs = {'img': out[0], 'big_img': out[1]} + # Model decides when action is completed, so desire input is just a pulse triggered on rising edge + inputs['desire_pulse'][0] = 0 + self.npy['desire'][:] = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) + self.prev_desire[:] = inputs['desire_pulse'] + self.npy['traffic_convention'][:] = inputs['traffic_convention'] + self.npy['tfm'][:,:] = transforms['img'][:,:] + self.npy['big_tfm'][:,:] = transforms['big_img'][:,:] if prepare_only: + self.warp_enqueue(**self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img']) return None - self.vision_output = self.vision_run(**vision_inputs).contiguous().realize().uop.base.buffer.numpy().flatten() - vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(self.vision_output, self.vision_output_slices)) + vision_output, policy_output = self.run_policy( + **self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'] + ) - self.full_input_queues.enqueue({'features_buffer': vision_outputs_dict['hidden_state'], 'desire_pulse': new_desire}) - for k in ['desire_pulse', 'features_buffer']: - self.numpy_inputs[k][:] = self.full_input_queues.get(k)[k] - self.numpy_inputs['traffic_convention'][:] = inputs['traffic_convention'] - if 'action_t' in self.numpy_inputs: - self.numpy_inputs['action_t'][:] = inputs['action_t'] + vision_output = vision_output.numpy().flatten() + policy_output = policy_output.numpy().flatten() + vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices)) + policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(policy_output, self.policy_output_slices)) + combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict} - self.policy_output = self.policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy().flatten() - policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(self.policy_output, self.policy_output_slices)) - - self.off_policy_output = self.off_policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy() - off_policy_outputs_dict = self.parser.parse_off_policy_outputs(self.slice_outputs(self.off_policy_output, self.off_policy_output_slices)) - - combined_outputs_dict = {**vision_outputs_dict, **off_policy_outputs_dict, **policy_outputs_dict} if SEND_RAW_PRED: - combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy(), self.off_policy_output.copy()]) - + combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), policy_output.copy()]) return combined_outputs_dict @@ -269,11 +157,6 @@ def main(demo=False): # also need to move the aux USB interrupts for good timings config_realtime_process(7, 54) - st = time.monotonic() - cloudlog.warning("loading model") - model = ModelState() - cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") - # visionipc clients while True: available_streams = VisionIpcClient.available_streams("camerad", block=False) @@ -297,6 +180,11 @@ def main(demo=False): if use_extra_client: cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})") + st = time.monotonic() + cloudlog.warning("loading model") + model = ModelState(vipc_client_main.width, vipc_client_main.height) + cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") + # messaging pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry"]) sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"]) @@ -399,14 +287,9 @@ def main(demo=False): 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} - frame_delay = DT_MDL # compensate for time passed since the frame was captured: current_time - timestamp_eof is 50ms on average - action_delay = DT_MDL / 2 # middle of the interval between model output (current state) and next frame (expected state) - lat_action_t = lat_delay + frame_delay + action_delay - long_action_t = long_delay + frame_delay + action_delay inputs:dict[str, np.ndarray] = { 'desire_pulse': vec_desire, 'traffic_convention': traffic_convention, - 'action_t': np.array([lat_action_t, long_action_t], dtype=np.float32), } mt1 = time.perf_counter() @@ -419,7 +302,9 @@ def main(demo=False): drivingdata_send = messaging.new_message('drivingModelData') posenet_send = messaging.new_message('cameraOdometry') - action = get_action_from_model(model_output, prev_action, lat_action_t, long_action_t, v_ego) + frame_delay = DT_MDL # compensate for time passed since the frame was captured: current_time - timestamp_eof is 50ms on average + action_delay = DT_MDL / 2 # middle of the interval between model output (current state) and next frame (expected state) + action = get_action_from_model(model_output, prev_action, lat_delay + frame_delay + action_delay, long_delay + frame_delay + action_delay, v_ego) prev_action = action fill_model_msg(drivingdata_send, modelv2_send, model_output, action, publish_state, meta_main.frame_id, meta_extra.frame_id, frame_id, diff --git a/selfdrive/modeld/models/big_driving_policy.onnx b/selfdrive/modeld/models/big_driving_policy.onnx new file mode 120000 index 00000000000000..e1b653a14a03d6 --- /dev/null +++ b/selfdrive/modeld/models/big_driving_policy.onnx @@ -0,0 +1 @@ +driving_policy.onnx \ No newline at end of file diff --git a/selfdrive/modeld/models/big_driving_vision.onnx b/selfdrive/modeld/models/big_driving_vision.onnx new file mode 120000 index 00000000000000..28ee71dd746e63 --- /dev/null +++ b/selfdrive/modeld/models/big_driving_vision.onnx @@ -0,0 +1 @@ +driving_vision.onnx \ No newline at end of file diff --git a/selfdrive/modeld/models/driving_off_policy.onnx b/selfdrive/modeld/models/driving_off_policy.onnx deleted file mode 100644 index 2975a571dcb392..00000000000000 --- a/selfdrive/modeld/models/driving_off_policy.onnx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9a2bdfb988ab96fc60d991fd9e6b38221819b235539bb0d6b781fc029438f599 -size 13928996 diff --git a/selfdrive/modeld/models/driving_on_policy.onnx b/selfdrive/modeld/models/driving_on_policy.onnx deleted file mode 100644 index 239249a81c8f82..00000000000000 --- a/selfdrive/modeld/models/driving_on_policy.onnx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0e66a9a3f5eb8da0b7693904ff8f1904b5fb43c9c0e265c170efad30c84629cb -size 12548985 diff --git a/selfdrive/modeld/models/driving_policy.onnx b/selfdrive/modeld/models/driving_policy.onnx new file mode 100644 index 00000000000000..611ae9fe85f837 --- /dev/null +++ b/selfdrive/modeld/models/driving_policy.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:78477124cbf3ffe30fa951ebada8410b43c4242c6054584d656f1d329b067e15 +size 14060847 diff --git a/selfdrive/modeld/models/driving_vision.onnx b/selfdrive/modeld/models/driving_vision.onnx index c01d6e233b8087..6c9fc4c84d3632 100644 --- a/selfdrive/modeld/models/driving_vision.onnx +++ b/selfdrive/modeld/models/driving_vision.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:5928713f355d75cf01ec4961bb0442c3712581dd3bdcf9aaca387eee77049f69 -size 23272727 +oid sha256:ee29ee5bce84d1ce23e9ff381280de9b4e4d96d2934cd751740354884e112c66 +size 46877473 diff --git a/selfdrive/modeld/parse_model_outputs.py b/selfdrive/modeld/parse_model_outputs.py index 3211668b5aa5cc..a0b45d2a981685 100644 --- a/selfdrive/modeld/parse_model_outputs.py +++ b/selfdrive/modeld/parse_model_outputs.py @@ -96,17 +96,11 @@ def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndar self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,)) self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) - self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH)) - self.parse_binary_crossentropy('meta', outs) - return outs - - def parse_off_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - plan_mhp = self.is_mhp(outs, 'plan', ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH) - plan_in_N, plan_out_N = (ModelConstants.PLAN_MHP_N, ModelConstants.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=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH)) self.parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_ROAD_EDGES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH)) self.parse_binary_crossentropy('lane_lines_prob', outs) + self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH)) + self.parse_binary_crossentropy('meta', outs) self.parse_binary_crossentropy('lead_prob', outs) lead_mhp = self.is_mhp(outs, 'lead', ModelConstants.LEAD_MHP_SELECTION * ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH) lead_in_N, lead_out_N = (ModelConstants.LEAD_MHP_N, ModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0) @@ -116,11 +110,15 @@ def parse_off_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np. return outs def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - self.parse_mdn('action', outs, in_N=0, out_N=0, out_shape=(ModelConstants.ACTION_WIDTH,)) + plan_mhp = self.is_mhp(outs, 'plan', ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH) + plan_in_N, plan_out_N = (ModelConstants.PLAN_MHP_N, ModelConstants.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=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) + if 'planplus' in outs: + self.parse_mdn('planplus', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) + self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,)) return outs def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: outs = self.parse_vision_outputs(outs) - outs = self.parse_off_policy_outputs(outs) outs = self.parse_policy_outputs(outs) return outs diff --git a/selfdrive/modeld/tinygrad_helpers.py b/selfdrive/modeld/tinygrad_helpers.py deleted file mode 100644 index 49a6ed6161855c..00000000000000 --- a/selfdrive/modeld/tinygrad_helpers.py +++ /dev/null @@ -1,12 +0,0 @@ -import json -import os -from pathlib import Path - -MODELS_DIR = Path(__file__).parent / 'models' -COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json' - - -def set_tinygrad_backend_from_compiled_flags() -> None: - if os.path.isfile(COMPILED_FLAGS_PATH): - with open(COMPILED_FLAGS_PATH) as f: - os.environ['DEV'] = str(json.load(f)['DEV']) From 66d9fe6dfd5d7989da881634cc5bdfcdf9548d85 Mon Sep 17 00:00:00 2001 From: Bruce Wayne Date: Mon, 4 May 2026 19:32:59 -0700 Subject: [PATCH 3/5] Add op11 models on PR37928 runtime --- selfdrive/modeld/SConscript | 9 ++-- selfdrive/modeld/compile_modeld.py | 52 ++++++++++++++----- selfdrive/modeld/constants.py | 1 + selfdrive/modeld/fill_model_msg.py | 5 +- selfdrive/modeld/modeld.py | 49 +++++++++-------- .../modeld/models/big_driving_policy.onnx | 1 - .../modeld/models/big_driving_vision.onnx | 1 - .../modeld/models/driving_off_policy.onnx | 3 ++ .../modeld/models/driving_on_policy.onnx | 3 ++ selfdrive/modeld/models/driving_policy.onnx | 3 -- selfdrive/modeld/models/driving_vision.onnx | 4 +- selfdrive/modeld/parse_model_outputs.py | 18 ++++--- 12 files changed, 95 insertions(+), 54 deletions(-) delete mode 120000 selfdrive/modeld/models/big_driving_policy.onnx delete mode 120000 selfdrive/modeld/models/big_driving_vision.onnx create mode 100644 selfdrive/modeld/models/driving_off_policy.onnx create mode 100644 selfdrive/modeld/models/driving_on_policy.onnx delete mode 100644 selfdrive/modeld/models/driving_policy.onnx diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 40b30b8c1bc4d0..80b764b27c20ce 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -56,7 +56,7 @@ compiled_flags_node = lenv.Command( mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else '' # Get model metadata -for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: +for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']: fn = File(f"models/{model_name}").abspath script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)] cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' @@ -65,8 +65,8 @@ for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: modeld_dir = Dir("#selfdrive/modeld").abspath compile_modeld_script = [File(f"{modeld_dir}/compile_modeld.py")] compile_dm_warp_script = [File(f"{modeld_dir}/compile_dm_warp.py")] -driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_policy']] -driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']] +driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_off_policy', 'driving_on_policy']] +driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_off_policy', 'driving_on_policy']] model_w, model_h = MEDMODEL_INPUT_SIZE frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ @@ -75,7 +75,8 @@ for cfg in MODELD_CONFIGS: f'--model-size {model_w}x{model_h} ' f'--nv12 {",".join(str(x) for x in cfg.nv12)} ' f'--vision-onnx {File("models/driving_vision.onnx").abspath} ' - f'--policy-onnx {File("models/driving_policy.onnx").abspath} ' + f'--off-policy-onnx {File("models/driving_off_policy.onnx").abspath} ' + f'--on-policy-onnx {File("models/driving_on_policy.onnx").abspath} ' f'--output {cfg.pkl_path} --frame-skip {frame_skip}' + (' --prepare-only' if cfg.prepare_only else '')) node = lenv.Command(cfg.pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + driving_metadata_deps + [chunker_file, compiled_flags_node], cmd) diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index 61de986d564f41..d7e5efaac1b070 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -94,6 +94,10 @@ def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip): 'tfm': np.zeros((3, 3), dtype=np.float32), 'big_tfm': np.zeros((3, 3), dtype=np.float32), } + if 'action_t' in policy_input_shapes: + npy['action_t'] = np.zeros(policy_input_shapes['action_t'], dtype=np.float32) + if 'prev_action' in policy_input_shapes: + npy['prev_action'] = np.zeros(policy_input_shapes['prev_action'][2], dtype=np.float32) input_queues = { 'img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(), 'big_img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(), @@ -101,6 +105,9 @@ def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip): 'desire_q': Tensor.zeros(frame_skip * dp[1], dp[0], dp[2]).contiguous().realize(), **{k: Tensor(v, device='NPY').realize() for k, v in npy.items()}, } + if 'prev_action' in policy_input_shapes: + pa = policy_input_shapes['prev_action'] # (1, 25, 2) + input_queues['prev_action_q'] = Tensor.zeros(frame_skip * (pa[1] - 1) + 1, pa[0], pa[2]).contiguous().realize() return input_queues, npy @@ -117,18 +124,24 @@ def sample_desire(buf, frame_skip): return buf.reshape(-1, frame_skip, *buf.shape[1:]).max(1).flatten(0, 1).unsqueeze(0) -def make_run_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, model_h, +def make_run_policy(vision_runner, off_policy_runner, on_policy_runner, nv12: NV12Frame, model_w, model_h, vision_features_slice, frame_skip, 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): + def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, action_t, tfm, big_tfm, frame, big_frame, + prev_action_q=None, prev_action=None): 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) + action_t = action_t.to(Device.DEFAULT) + to_realize = [tfm, big_tfm, desire, traffic_convention, action_t] + if prev_action is not None: + prev_action = prev_action.to(Device.DEFAULT) + to_realize.append(prev_action) + Tensor.realize(*to_realize) 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) @@ -142,30 +155,42 @@ def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, t 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_pulse': desire_buf, 'traffic_convention': traffic_convention} - policy_out = next(iter(policy_runner(inputs).values())).cast('float32') - - return vision_out, policy_out + inputs = { + 'features_buffer': feat_buf, + 'desire_pulse': desire_buf, + 'traffic_convention': traffic_convention, + 'action_t': action_t, + } + if prev_action_q is not None and prev_action is not None: + inputs['prev_action'] = shift_and_sample(prev_action_q, prev_action.reshape(1, 1, -1), sample_skip_fn) + on_policy_out = next(iter(on_policy_runner(inputs).values())).cast('float32') + off_policy_out = next(iter(off_policy_runner(inputs).values())).cast('float32') + + return vision_out, on_policy_out, off_policy_out return run_policy def compile_modeld(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, - vision_onnx, policy_onnx, pkl_path): + vision_onnx, off_policy_onnx, on_policy_onnx, pkl_path): from get_model_metadata import metadata_path_for print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") vision_runner = OnnxRunner(vision_onnx) - policy_runner = OnnxRunner(policy_onnx) + off_policy_runner = OnnxRunner(off_policy_onnx) + on_policy_runner = OnnxRunner(on_policy_onnx) with open(metadata_path_for(vision_onnx), 'rb') as f: vision_metadata = pickle.load(f) vision_features_slice = vision_metadata['output_slices']['hidden_state'] vision_input_shapes = vision_metadata['input_shapes'] - with open(metadata_path_for(policy_onnx), 'rb') as f: + with open(metadata_path_for(on_policy_onnx), 'rb') as f: policy_input_shapes = pickle.load(f)['input_shapes'] + with open(metadata_path_for(off_policy_onnx), 'rb') as f: + off_policy_input_shapes = pickle.load(f)['input_shapes'] + assert policy_input_shapes == off_policy_input_shapes - _run = make_run_policy(vision_runner, policy_runner, nv12, model_w, model_h, + _run = make_run_policy(vision_runner, off_policy_runner, on_policy_runner, nv12, model_w, model_h, vision_features_slice, frame_skip, prepare_only) run_policy_jit = TinyJit(_run, prune=True) @@ -235,7 +260,8 @@ def _parse_nv12(s): p.add_argument('--nv12', type=_parse_nv12, required=True, help=f'NV12 frame layout: {",".join(NV12Frame._fields)}') p.add_argument('--vision-onnx', required=True) - p.add_argument('--policy-onnx', required=True) + p.add_argument('--off-policy-onnx', required=True) + p.add_argument('--on-policy-onnx', required=True) p.add_argument('--output', required=True) p.add_argument('--prepare-only', action='store_true') p.add_argument('--frame-skip', type=int, required=True) @@ -243,4 +269,4 @@ def _parse_nv12(s): model_w, model_h = args.model_size compile_modeld(args.nv12, model_w, model_h, args.prepare_only, args.frame_skip, - args.vision_onnx, args.policy_onnx, args.output) + args.vision_onnx, args.off_policy_onnx, args.on_policy_onnx, args.output) diff --git a/selfdrive/modeld/constants.py b/selfdrive/modeld/constants.py index ff7e1d86006e83..0fb09262d0192e 100644 --- a/selfdrive/modeld/constants.py +++ b/selfdrive/modeld/constants.py @@ -38,6 +38,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/selfdrive/modeld/fill_model_msg.py b/selfdrive/modeld/fill_model_msg.py index 82c4c92b1d53c7..92a2dfa58d7f3a 100644 --- a/selfdrive/modeld/fill_model_msg.py +++ b/selfdrive/modeld/fill_model_msg.py @@ -125,7 +125,10 @@ def fill_model_msg(base_msg: capnp._DynamicStructBuilder, extended_msg: capnp._D # meta meta = modelV2.meta - meta.desireState = net_output_data['desire_state'][0].reshape(-1).tolist() + if 'desire_state' in net_output_data: + meta.desireState = net_output_data['desire_state'][0].reshape(-1).tolist() + else: + meta.desireState = [0.0] * ModelConstants.DESIRE_PRED_WIDTH meta.desirePrediction = net_output_data['desire_pred'][0].reshape(-1).tolist() meta.engagedProb = net_output_data['meta'][0,Meta.ENGAGED].item() meta.init('disengagePredictions') diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 73ed19ec943790..7041d860e382e6 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -36,7 +36,8 @@ SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl' -POLICY_METADATA_PATH = MODELS_DIR / 'driving_policy_metadata.pkl' +OFF_POLICY_METADATA_PATH = MODELS_DIR / 'driving_off_policy_metadata.pkl' +ON_POLICY_METADATA_PATH = MODELS_DIR / 'driving_on_policy_metadata.pkl' LAT_SMOOTH_SECONDS = 0.0 LONG_SMOOTH_SECONDS = 0.3 @@ -44,20 +45,12 @@ -def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, - lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action: - plan = model_output['plan'][0] - desired_accel, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0], - plan[:,Plan.ACCELERATION][:,0], - ModelConstants.T_IDXS, - action_t=long_action_t) - desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, LONG_SMOOTH_SECONDS) +def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, v_ego: float) -> log.ModelDataV2.Action: + desired_curv_unscaled, desired_accel = model_output['action'][0] + desired_curvature = desired_curv_unscaled / 100 + should_stop = (v_ego < 0.3 and desired_accel < 0.1) - desired_curvature = get_curvature_from_plan(plan[:,Plan.T_FROM_CURRENT_EULER][:,2], - plan[:,Plan.ORIENTATION_RATE][:,2], - ModelConstants.T_IDXS, - v_ego, - lat_action_t) + desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, LONG_SMOOTH_SECONDS) if v_ego > MIN_LAT_CONTROL_SPEED: desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, LAT_SMOOTH_SECONDS) else: @@ -87,7 +80,11 @@ def __init__(self, cam_w: int, cam_h: int): self.vision_input_names = list(self.vision_input_shapes.keys()) self.vision_output_slices = vision_metadata['output_slices'] - with open(POLICY_METADATA_PATH, 'rb') as f: + with open(OFF_POLICY_METADATA_PATH, 'rb') as f: + off_policy_metadata = pickle.load(f) + self.off_policy_output_slices = off_policy_metadata['output_slices'] + + with open(ON_POLICY_METADATA_PATH, 'rb') as f: policy_metadata = pickle.load(f) self.policy_input_shapes = policy_metadata['input_shapes'] self.policy_output_slices = policy_metadata['output_slices'] @@ -127,6 +124,10 @@ def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], self.npy['desire'][:] = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) self.prev_desire[:] = inputs['desire_pulse'] self.npy['traffic_convention'][:] = inputs['traffic_convention'] + if 'action_t' in self.npy: + self.npy['action_t'][:] = inputs['action_t'] + if 'prev_action' in self.npy: + self.npy['prev_action'][:] = inputs['prev_action'] self.npy['tfm'][:,:] = transforms['img'][:,:] self.npy['big_tfm'][:,:] = transforms['big_img'][:,:] @@ -134,18 +135,20 @@ def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], self.warp_enqueue(**self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img']) return None - vision_output, policy_output = self.run_policy( + vision_output, policy_output, off_policy_output = self.run_policy( **self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'] ) vision_output = vision_output.numpy().flatten() + off_policy_output = off_policy_output.numpy().flatten() policy_output = policy_output.numpy().flatten() vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(vision_output, self.vision_output_slices)) + off_policy_outputs_dict = self.parser.parse_off_policy_outputs(self.slice_outputs(off_policy_output, self.off_policy_output_slices)) policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(policy_output, self.policy_output_slices)) - combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict} + combined_outputs_dict = {**vision_outputs_dict, **off_policy_outputs_dict, **policy_outputs_dict} if SEND_RAW_PRED: - combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), policy_output.copy()]) + combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), policy_output.copy(), off_policy_output.copy()]) return combined_outputs_dict @@ -287,9 +290,15 @@ def main(demo=False): 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} + frame_delay = DT_MDL # compensate for time passed since the frame was captured: current_time - timestamp_eof is 50ms on average + action_delay = DT_MDL / 2 # middle of the interval between model output (current state) and next frame (expected state) + lat_action_t = lat_delay + frame_delay + action_delay + long_action_t = long_delay + frame_delay + action_delay inputs:dict[str, np.ndarray] = { 'desire_pulse': vec_desire, 'traffic_convention': traffic_convention, + 'action_t': np.array([lat_action_t, long_action_t], dtype=np.float32), + 'prev_action': np.array([prev_action.desiredCurvature * max(1.0, v_ego)**2, prev_action.desiredAcceleration], dtype=np.float32), } mt1 = time.perf_counter() @@ -302,9 +311,7 @@ def main(demo=False): drivingdata_send = messaging.new_message('drivingModelData') posenet_send = messaging.new_message('cameraOdometry') - frame_delay = DT_MDL # compensate for time passed since the frame was captured: current_time - timestamp_eof is 50ms on average - action_delay = DT_MDL / 2 # middle of the interval between model output (current state) and next frame (expected state) - action = get_action_from_model(model_output, prev_action, lat_delay + frame_delay + action_delay, long_delay + frame_delay + action_delay, v_ego) + action = get_action_from_model(model_output, prev_action, v_ego) prev_action = action fill_model_msg(drivingdata_send, modelv2_send, model_output, action, publish_state, meta_main.frame_id, meta_extra.frame_id, frame_id, diff --git a/selfdrive/modeld/models/big_driving_policy.onnx b/selfdrive/modeld/models/big_driving_policy.onnx deleted file mode 120000 index e1b653a14a03d6..00000000000000 --- a/selfdrive/modeld/models/big_driving_policy.onnx +++ /dev/null @@ -1 +0,0 @@ -driving_policy.onnx \ No newline at end of file diff --git a/selfdrive/modeld/models/big_driving_vision.onnx b/selfdrive/modeld/models/big_driving_vision.onnx deleted file mode 120000 index 28ee71dd746e63..00000000000000 --- a/selfdrive/modeld/models/big_driving_vision.onnx +++ /dev/null @@ -1 +0,0 @@ -driving_vision.onnx \ No newline at end of file diff --git a/selfdrive/modeld/models/driving_off_policy.onnx b/selfdrive/modeld/models/driving_off_policy.onnx new file mode 100644 index 00000000000000..2975a571dcb392 --- /dev/null +++ b/selfdrive/modeld/models/driving_off_policy.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9a2bdfb988ab96fc60d991fd9e6b38221819b235539bb0d6b781fc029438f599 +size 13928996 diff --git a/selfdrive/modeld/models/driving_on_policy.onnx b/selfdrive/modeld/models/driving_on_policy.onnx new file mode 100644 index 00000000000000..239249a81c8f82 --- /dev/null +++ b/selfdrive/modeld/models/driving_on_policy.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0e66a9a3f5eb8da0b7693904ff8f1904b5fb43c9c0e265c170efad30c84629cb +size 12548985 diff --git a/selfdrive/modeld/models/driving_policy.onnx b/selfdrive/modeld/models/driving_policy.onnx deleted file mode 100644 index 611ae9fe85f837..00000000000000 --- a/selfdrive/modeld/models/driving_policy.onnx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:78477124cbf3ffe30fa951ebada8410b43c4242c6054584d656f1d329b067e15 -size 14060847 diff --git a/selfdrive/modeld/models/driving_vision.onnx b/selfdrive/modeld/models/driving_vision.onnx index 6c9fc4c84d3632..c01d6e233b8087 100644 --- a/selfdrive/modeld/models/driving_vision.onnx +++ b/selfdrive/modeld/models/driving_vision.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:ee29ee5bce84d1ce23e9ff381280de9b4e4d96d2934cd751740354884e112c66 -size 46877473 +oid sha256:5928713f355d75cf01ec4961bb0442c3712581dd3bdcf9aaca387eee77049f69 +size 23272727 diff --git a/selfdrive/modeld/parse_model_outputs.py b/selfdrive/modeld/parse_model_outputs.py index a0b45d2a981685..3211668b5aa5cc 100644 --- a/selfdrive/modeld/parse_model_outputs.py +++ b/selfdrive/modeld/parse_model_outputs.py @@ -96,11 +96,17 @@ def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndar self.parse_mdn('pose', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) self.parse_mdn('wide_from_device_euler', outs, in_N=0, out_N=0, out_shape=(ModelConstants.WIDE_FROM_DEVICE_WIDTH,)) self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,)) + self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH)) + self.parse_binary_crossentropy('meta', outs) + return outs + + def parse_off_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + plan_mhp = self.is_mhp(outs, 'plan', ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH) + plan_in_N, plan_out_N = (ModelConstants.PLAN_MHP_N, ModelConstants.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=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH)) self.parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_ROAD_EDGES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH)) self.parse_binary_crossentropy('lane_lines_prob', outs) - self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH)) - self.parse_binary_crossentropy('meta', outs) self.parse_binary_crossentropy('lead_prob', outs) lead_mhp = self.is_mhp(outs, 'lead', ModelConstants.LEAD_MHP_SELECTION * ModelConstants.LEAD_TRAJ_LEN * ModelConstants.LEAD_WIDTH) lead_in_N, lead_out_N = (ModelConstants.LEAD_MHP_N, ModelConstants.LEAD_MHP_SELECTION) if lead_mhp else (0, 0) @@ -110,15 +116,11 @@ def parse_vision_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndar return outs def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - plan_mhp = self.is_mhp(outs, 'plan', ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH) - plan_in_N, plan_out_N = (ModelConstants.PLAN_MHP_N, ModelConstants.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=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) - if 'planplus' in outs: - self.parse_mdn('planplus', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) - self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,)) + self.parse_mdn('action', outs, in_N=0, out_N=0, out_shape=(ModelConstants.ACTION_WIDTH,)) return outs def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: outs = self.parse_vision_outputs(outs) + outs = self.parse_off_policy_outputs(outs) outs = self.parse_policy_outputs(outs) return outs From ece8f7664d366eb400fd3c763bd427b3377e19e7 Mon Sep 17 00:00:00 2001 From: Bruce Wayne Date: Tue, 5 May 2026 09:32:09 -0700 Subject: [PATCH 4/5] Export off-policy desire state --- selfdrive/modeld/models/driving_off_policy.onnx | 4 ++-- selfdrive/modeld/parse_model_outputs.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/selfdrive/modeld/models/driving_off_policy.onnx b/selfdrive/modeld/models/driving_off_policy.onnx index 2975a571dcb392..d83ec7f1372d36 100644 --- a/selfdrive/modeld/models/driving_off_policy.onnx +++ b/selfdrive/modeld/models/driving_off_policy.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:9a2bdfb988ab96fc60d991fd9e6b38221819b235539bb0d6b781fc029438f599 -size 13928996 +oid sha256:a1facb2c1028a945b44179c3dfe87aea6170446eef7ec1bde177f8bc0ae5c2be +size 13981538 diff --git a/selfdrive/modeld/parse_model_outputs.py b/selfdrive/modeld/parse_model_outputs.py index 3211668b5aa5cc..1a7d699e687ca7 100644 --- a/selfdrive/modeld/parse_model_outputs.py +++ b/selfdrive/modeld/parse_model_outputs.py @@ -113,6 +113,7 @@ def parse_off_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np. lead_out_shape = (ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) if lead_mhp else \ (ModelConstants.LEAD_MHP_SELECTION, ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH) self.parse_mdn('lead', outs, in_N=lead_in_N, out_N=lead_out_N, out_shape=lead_out_shape) + self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,)) return outs def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: From 9757bf109e0673eb7481246d7dad91b3a70b51e9 Mon Sep 17 00:00:00 2001 From: Bruce Wayne Date: Tue, 5 May 2026 19:17:59 -0700 Subject: [PATCH 5/5] c93fb62d-9de8-4ebf-8529-9784feac48a7/100 --- selfdrive/modeld/models/driving_off_policy.onnx | 4 ++-- selfdrive/modeld/models/driving_on_policy.onnx | 4 ++-- selfdrive/modeld/models/driving_vision.onnx | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/selfdrive/modeld/models/driving_off_policy.onnx b/selfdrive/modeld/models/driving_off_policy.onnx index d83ec7f1372d36..51b659a08ea6cd 100644 --- a/selfdrive/modeld/models/driving_off_policy.onnx +++ b/selfdrive/modeld/models/driving_off_policy.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:a1facb2c1028a945b44179c3dfe87aea6170446eef7ec1bde177f8bc0ae5c2be -size 13981538 +oid sha256:dff39ef4705f7dc092ea6a867d030c1092d40fa124f15c36ed01ee62877a1f21 +size 18178659 diff --git a/selfdrive/modeld/models/driving_on_policy.onnx b/selfdrive/modeld/models/driving_on_policy.onnx index 239249a81c8f82..50ffcfec4cae96 100644 --- a/selfdrive/modeld/models/driving_on_policy.onnx +++ b/selfdrive/modeld/models/driving_on_policy.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:0e66a9a3f5eb8da0b7693904ff8f1904b5fb43c9c0e265c170efad30c84629cb -size 12548985 +oid sha256:4e7e31a20e93c6e0fd37d29c339e6668a915e3550d971a17e6560b9227b7ffcd +size 16745129 diff --git a/selfdrive/modeld/models/driving_vision.onnx b/selfdrive/modeld/models/driving_vision.onnx index c01d6e233b8087..33601fc1ff2121 100644 --- a/selfdrive/modeld/models/driving_vision.onnx +++ b/selfdrive/modeld/models/driving_vision.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:5928713f355d75cf01ec4961bb0442c3712581dd3bdcf9aaca387eee77049f69 -size 23272727 +oid sha256:50070f898ad223475fced42d12a8e67e57e7fac7b21ddc26b05e64f1ae536880 +size 23210375