diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index 8d6058f39..81d0a3003 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import argparse import atexit +import math import os import pickle import time @@ -33,7 +34,7 @@ from tinygrad.engine.jit import TinyJit NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size']) WARP_INPUTS = ['img_q', 'big_img_q', 'tfm', 'big_tfm'] -POLICY_INPUTS = ['feat_q', 'desire_q', 'desire', 'traffic_convention', 'action_t'] +POLICY_INPUTS = ['feat_q', 'desire_q', 'packed_npy_inputs'] 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) @@ -130,25 +131,28 @@ def make_warp_input_queues(vision_input_shapes, frame_skip, device): return input_queues, npy +def get_policy_npy_shapes(policy_input_shapes): + 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 + shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(tc)} + return shapes, [math.prod(s) for s in shapes.values()] + + 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 - policy_npy = { - 'desire': np.zeros(dp[2], dtype=np.float32), - 'traffic_convention': np.zeros(tc, dtype=np.float32), - 'action_t': np.zeros(at, dtype=np.float32), - } - npy.update(policy_npy) + shapes, sizes = get_policy_npy_shapes(policy_input_shapes) + packed_npy_inputs = np.zeros(sum(sizes), dtype=np.float32) + # views into the packed inputs, to be refilled at runtime + npy.update({k: v.reshape(s) for (k, s), v in zip(shapes.items(), np.split(packed_npy_inputs, np.cumsum(sizes[:-1])), strict=True)}) 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 policy_npy.items()}, + 'packed_npy_inputs': Tensor(packed_npy_inputs, device='NPY').realize(), }) return input_queues, npy @@ -175,10 +179,11 @@ def make_warp(nv12, model_w, model_h, frame_skip): big_tfm = big_tfm.to(WARP_DEV) Tensor.realize(tfm, big_tfm) - warped_frame = frame_prepare(frame, tfm).unsqueeze(0).to(Device.DEFAULT) - warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0).to(Device.DEFAULT) - img = shift_and_sample(img_q, warped_frame, sample_skip_fn) - big_img = shift_and_sample(big_img_q, warped_big_frame, sample_skip_fn) + warped_frame = frame_prepare(frame, tfm).unsqueeze(0) + warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0) + warped = Tensor.cat(warped_frame, warped_big_frame).to(Device.DEFAULT) + img = shift_and_sample(img_q, warped[0:1], sample_skip_fn) + big_img = shift_and_sample(big_img_q, warped[1:2], sample_skip_fn) return img, big_img return warp_enqueue @@ -187,12 +192,11 @@ 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'] + npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['on_policy']['input_shapes']) - def run_policy(img, big_img, feat_q, desire_q, desire, traffic_convention, action_t): - desire = desire.to(Device.DEFAULT) - traffic_convention = traffic_convention.to(Device.DEFAULT) - action_t = action_t.to(Device.DEFAULT) - Tensor.realize(desire, traffic_convention, action_t) + def run_policy(img, big_img, feat_q, desire_q, packed_npy_inputs): + packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT).realize() + desire, traffic_convention, action_t = (t.reshape(s) for t, s in zip(packed_npy_inputs.split(npy_sizes), npy_shapes.values(), strict=True)) desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn) vision_out = next(iter(model_runners['vision']({'img': img, 'big_img': big_img}).values())).cast('float32') @@ -207,7 +211,7 @@ def make_run_policy(model_runners, model_metadata, frame_skip): } on_policy_out = next(iter(model_runners['on_policy'](inputs).values())).cast('float32') off_policy_out = next(iter(model_runners['off_policy'](inputs).values())).cast('float32') - return vision_out, on_policy_out, off_policy_out + return Tensor.cat(vision_out, on_policy_out, off_policy_out, dim=1), return run_policy diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index cf53eb23d..7e4976bf5 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -83,6 +83,7 @@ class ModelState: 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'] + self.vision_output_len = vision_metadata['output_shapes']['outputs'][1] off_policy_metadata = jits['metadata']['off_policy'] self.off_policy_output_slices = off_policy_metadata['output_slices'] @@ -90,6 +91,7 @@ class ModelState: policy_metadata = jits['metadata']['on_policy'] self.policy_input_shapes = policy_metadata['input_shapes'] self.policy_output_slices = policy_metadata['output_slices'] + self.on_policy_output_len = policy_metadata['output_shapes']['outputs'][1] self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32) @@ -131,13 +133,11 @@ class ModelState: if prepare_only: return None - vision_output, on_policy_output, off_policy_output = self.run_policy( + outs, = self.run_policy( **{k: self.input_queues[k] for k in POLICY_INPUTS if k in self.input_queues}, img=img, big_img=big_img ) + vision_output, on_policy_output, off_policy_output = np.split(outs.numpy()[0], [self.vision_output_len, self.vision_output_len+self.on_policy_output_len]) - vision_output = vision_output.numpy().flatten() - off_policy_output = off_policy_output.numpy().flatten() - on_policy_output = on_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(on_policy_output, self.policy_output_slices))