mawei/sunnypilot
1
0
Fork 0
mirror of https://github.com/sunnypilot/sunnypilot.git synced 2026-07-16 15:12:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor compile_modeld model setup (#38128)

Browse Source
This commit is contained in:
Harald Schäfer
2026-06-02 16:35:40 -07:00
committed by GitHub
parent 69e2c321e4
commit d09a411cb4
1 changed files with 42 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
selfdrive/modeld/compile_modeld.py
+42 -29
View File
@@ -5,7 +5,7 @@ import os
import pickle
import time
from functools import partial
from collections import namedtuple, defaultdict
from collections import namedtuple
import numpy as np
@@ -113,31 +113,43 @@ def make_frame_prepare(nv12: NV12Frame, model_w, model_h):
return frame_prepare_tinygrad
def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device):
def make_warp_input_queues(vision_input_shapes, frame_skip, device):
img = vision_input_shapes['img'] # (1, 12, 128, 256)
n_frames = img[1] // 6
img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3])
npy = {
'tfm': np.zeros((3, 3), dtype=np.float32),
'big_tfm': np.zeros((3, 3), dtype=np.float32),
}
input_queues = {
'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
**{k: Tensor(v, device='NPY').realize() for k, v in npy.items()},
}
return input_queues, npy
def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, device):
input_queues, npy = make_warp_input_queues(vision_input_shapes, frame_skip, device)
fb = policy_input_shapes['features_buffer'] # (1, 25, 512)
dp = policy_input_shapes['desire_pulse'] # (1, 25, 8)
tc = policy_input_shapes['traffic_convention'] # (1, 2)
#TODO action_t is hardcoded to match tc for future compatibility
at = tc
npy = {
policy_npy = {
'desire': np.zeros(dp[2], dtype=np.float32),
'traffic_convention': np.zeros(tc, dtype=np.float32),
'tfm': np.zeros((3, 3), dtype=np.float32),
'big_tfm': np.zeros((3, 3), dtype=np.float32),
'action_t': np.zeros(at, dtype=np.float32),
}
input_queues = {
'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
npy.update(policy_npy)
input_queues.update({
'feat_q': Tensor(np.zeros((frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]), dtype=np.float32), device=device).contiguous().realize(),
'desire_q': Tensor(np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32), device=device).contiguous().realize(),
**{k: Tensor(v, device='NPY').realize() for k, v in npy.items()},
}
**{k: Tensor(v, device='NPY').realize() for k, v in policy_npy.items()},
})
return input_queues, npy
@@ -171,9 +183,10 @@ def make_warp(nv12, model_w, model_h, frame_skip):
return warp_enqueue
def make_run_policy(vision_runner, on_policy_runner, vision_features_slice, frame_skip):
def make_run_policy(model_runners, model_metadata, frame_skip):
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
sample_skip_fn = partial(sample_skip, frame_skip=frame_skip)
vision_features_slice = model_metadata['vision']['output_slices']['hidden_state']
def run_policy(img, big_img, feat_q, desire_q, desire, traffic_convention, action_t):
desire = desire.to(Device.DEFAULT)
@@ -181,7 +194,7 @@ def make_run_policy(vision_runner, on_policy_runner, vision_features_slice, fram
action_t = action_t.to(Device.DEFAULT)
Tensor.realize(desire, traffic_convention, action_t)
desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
vision_out = next(iter(vision_runner({'img': img, 'big_img': big_img}).values())).cast('float32')
vision_out = next(iter(model_runners['vision']({'img': img, 'big_img': big_img}).values())).cast('float32')
new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
@@ -192,20 +205,16 @@ def make_run_policy(vision_runner, on_policy_runner, vision_features_slice, fram
'traffic_convention': traffic_convention,
'action_t': action_t,
}
on_policy_out = next(iter(on_policy_runner(inputs).values())).cast('float32')
on_policy_out = next(iter(model_runners['on_policy'](inputs).values())).cast('float32')
#off_policy_out = next(iter(off_policy_runner(inputs).values())).cast('float32')
return vision_out, on_policy_out
return run_policy
def compile_jit(jit, make_random_inputs, input_keys, frame_skip, vision_metadata, policy_metadata):
vision_input_shapes = vision_metadata['input_shapes']
policy_input_shapes = policy_metadata['input_shapes']
def compile_jit(jit, make_random_inputs, input_keys, make_queues):
SEED = 42
def random_inputs_run(fn, seed, test_val=None, test_buffers=None, expect_match=True):
input_queues, npy = make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip, Device.DEFAULT)
input_queues, npy = make_queues(Device.DEFAULT)
np.random.seed(seed)
Tensor.manual_seed(seed)
@@ -274,25 +283,29 @@ if __name__ == "__main__":
p.add_argument('--frame-skip', type=int, required=True)
args = p.parse_args()
out = defaultdict(dict)
vision_path, on_policy_path = read_file_chunked_to_shm(args.vision_onnx), read_file_chunked_to_shm(args.on_policy_onnx)
model_paths = {
'vision': read_file_chunked_to_shm(args.vision_onnx),
'on_policy': read_file_chunked_to_shm(args.on_policy_onnx),
}
model_w, model_h = args.model_size
vision_runner = OnnxRunner(vision_path)
on_policy_runner = OnnxRunner(on_policy_path)
vision_metadata, on_policy_metadata = make_metadata_dict(vision_path), make_metadata_dict(on_policy_path)
model_runners = {name: OnnxRunner(path) for name, path in model_paths.items()}
out = {'metadata': {name: make_metadata_dict(path) for name, path in model_paths.items()}}
run_policy_jit = TinyJit(make_run_policy(vision_runner, on_policy_runner, vision_metadata['output_slices']['hidden_state'], args.frame_skip), prune=True)
out['metadata']['vision'], out['metadata']['on_policy'] = vision_metadata, on_policy_metadata
run_policy_jit = TinyJit(make_run_policy(model_runners, out['metadata'], args.frame_skip), prune=True)
make_random_model_inputs = partial(make_random_images, keys=['img', 'big_img'], shape=vision_metadata['input_shapes']['img'])
out['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS, args.frame_skip, vision_metadata, on_policy_metadata)
make_policy_queues = partial(make_input_queues, out['metadata']['vision']['input_shapes'],
out['metadata']['on_policy']['input_shapes'], args.frame_skip)
make_random_model_inputs = partial(make_random_images, keys=['img', 'big_img'], shape=out['metadata']['vision']['input_shapes']['img'])
out['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS,
make_policy_queues)
for cam_w, cam_h in args.camera_resolutions:
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=WARP_DEV)
warp_enqueue = TinyJit(make_warp(nv12, model_w, model_h, args.frame_skip), prune=True)
out[(cam_w,cam_h)] = compile_jit(warp_enqueue, make_random_warp_inputs, WARP_INPUTS, args.frame_skip, vision_metadata, on_policy_metadata)
make_warp_queues = partial(make_warp_input_queues, out['metadata']['vision']['input_shapes'], args.frame_skip)
out[(cam_w,cam_h)] = compile_jit(warp_enqueue, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)
with open(args.output, "wb") as f:
pickle.dump(out, f)