diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index 2f91076ab7..a7298705ef 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -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)