diff --git a/openpilot/selfdrive/modeld/SConscript b/openpilot/selfdrive/modeld/SConscript index 2f3e19a315..a302016460 100644 --- a/openpilot/selfdrive/modeld/SConscript +++ b/openpilot/selfdrive/modeld/SConscript @@ -6,7 +6,6 @@ from SCons.Script import Action, Value from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks 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 TG_INPUT_DEVICES_PATH, chestnut_present, modeld_pkl_path @@ -68,7 +67,6 @@ compile_modeld_script = [ File("#openpilot/common/hardware/hw.py"), ] model_w, model_h = MEDMODEL_INPUT_SIZE -frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ for chestnut in [False, True] if CHESTNUT else [False]: target_pkl_path = File(modeld_pkl_path(chestnut)).abspath @@ -81,7 +79,7 @@ for chestnut in [False, True] if CHESTNUT else [False]: f'--model-size {model_w}x{model_h} ' f'--camera-resolutions {camera_res_args} ' f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} ' - f'--output {target_pkl_path} --frame-skip {frame_skip}') + f'--output {target_pkl_path}') onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps) chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum * len(camera_configs))) def do_compile(target, source, env, command=cmd, pkl=target_pkl_path, chunks=chunk_targets): diff --git a/openpilot/selfdrive/modeld/compile_modeld.py b/openpilot/selfdrive/modeld/compile_modeld.py index 52be6897c0..c2fd116f1a 100755 --- a/openpilot/selfdrive/modeld/compile_modeld.py +++ b/openpilot/selfdrive/modeld/compile_modeld.py @@ -37,7 +37,6 @@ from tinygrad.engine.jit import TinyJit NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size']) -MODELD_INPUTS = ['img_q', 'big_img_q', 'feat_q', 'desire_q', 'packed_npy_inputs'] def nv12_copy_size(stride: int, y_height: int, uv_height: int) -> int: @@ -113,58 +112,26 @@ def make_frame_prepare(nv12: NV12Frame, model_w, model_h): return frame_prepare_tinygrad -def get_policy_npy_shapes(input_shapes): - dp = input_shapes['desire_pulse'] # (1, 25, 8) - tc = input_shapes['traffic_convention'] # (1, 2) - at = input_shapes['action_t'] # (1, 2) - fb = input_shapes['features_buffer'] # (1, T-1, ...) e.g. (1, 24, 32, 512) with spatial features - feat_dim = math.prod(fb[2:]) - # TODO prev_feat shouldn't exist and be handled inside the JIT, but corrupt on QCOM for now - shapes = {'desire': (dp[2],), 'traffic_convention': tuple(tc), 'action_t': tuple(at), 'prev_feat': (fb[0], feat_dim)} +def get_npy_shapes(input_shapes, state_pairs): + shapes = {'tfm': (3, 3), 'big_tfm': (3, 3)} | { + name: shape for name, (shape, _) in input_shapes.items() if name not in state_pairs and name != 'new_img'} return shapes, [math.prod(s) for s in shapes.values()] -def make_input_queues(input_shapes, frame_skip, device, frame_copy_size): - img = input_shapes['img'] # (1, 12, 128, 256) - fb = input_shapes['features_buffer'] # (1, T-1, ...), past features only; the model appends the current frame's feature - feat_dim = math.prod(fb[2:]) - dp = input_shapes['desire_pulse'] # (1, 25, 8) - n_frames = img[1] // 6 - img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3]) - - policy_shapes, _ = get_policy_npy_shapes(input_shapes) - shapes = {'tfm': (3, 3), 'big_tfm': (3, 3)} | policy_shapes - sizes = [math.prod(s) for s in shapes.values()] +def make_input_queues(input_shapes, state_pairs, device, frame_copy_size): + shapes, sizes = get_npy_shapes(input_shapes, state_pairs) packed_npy_size = sum(sizes) * np.dtype(np.float32).itemsize packed_input = np.zeros(packed_npy_size + 2 * frame_copy_size, dtype=np.uint8) packed_npy_inputs = packed_input[:packed_npy_size].view(np.float32) frames = packed_input[packed_npy_size:] frame_views = {'img': frames[:frame_copy_size], 'big_img': frames[frame_copy_size:]} - # views into the packed inputs, to be refilled at runtime npy = {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 = { - 'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), - 'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(), - 'feat_q': Tensor(np.zeros((frame_skip * fb[1], fb[0], feat_dim), 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(), - 'packed_npy_inputs': Tensor(packed_input, device='NPY').realize(), - } + input_queues = {name: Tensor(np.zeros(shape, dtype=dtype.fmt), device=device).realize() + for name, (shape, dtype) in input_shapes.items() if name in state_pairs} + input_queues['packed_npy_inputs'] = Tensor(packed_input, device='NPY').realize() return input_queues, npy, frame_views -def shift_and_sample(buf, new_val, sample_fn): - buf.assign(buf[1:].cat(new_val, dim=0).contiguous()) - return sample_fn(buf) - - -def sample_skip(buf, frame_skip): - return buf[::frame_skip].contiguous().flatten(0, 1).unsqueeze(0) - - -def sample_desire(buf, frame_skip): - return buf.reshape(-1, frame_skip, *buf.shape[1:]).max(1).flatten(0, 1).unsqueeze(0) - - def make_warp(nv12, model_w, model_h): frame_prepare = make_frame_prepare(nv12, model_w, model_h) @@ -182,54 +149,27 @@ def make_warp(nv12, model_w, model_h): return warp -def make_run_policy(model_runner, model_metadata, frame_skip): - sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) - sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) - npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['input_shapes']) - model_input_dtypes = {name: spec.dtype for name, spec in model_runner.graph_inputs.items()} +def make_run_model(warp, model_runner, input_shapes, state_pairs, frame_copy_size): + shapes, sizes = get_npy_shapes(input_shapes, state_pairs) + packed_npy_size = sum(sizes) * np.dtype(np.float32).itemsize - def run_policy(warped, img_q, big_img_q, feat_q, desire_q, packed_npy_inputs): - packed_npy_inputs = packed_npy_inputs.to(Device.DEFAULT) - Tensor.realize(packed_npy_inputs, warped) - - 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) - - desire, traffic_convention, action_t, prev_feat = (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) - feat_buf = shift_and_sample(feat_q, prev_feat.reshape(1, 1, -1), sample_skip_fn) - - inputs = { - 'img': img, - 'big_img': big_img, - 'features_buffer': feat_buf.reshape(model_metadata['input_shapes']['features_buffer']), - 'desire_pulse': desire_buf, - 'traffic_convention': traffic_convention, - 'action_t': action_t, - } - inputs = {name: value.cast(model_input_dtypes[name]) for name, value in inputs.items()} - out = next(iter(model_runner(inputs).values())).cast('float32') - return out, - return run_policy - - -def make_run_model(warp, run_policy, model_metadata, frame_copy_size): - _, policy_sizes = get_policy_npy_shapes(model_metadata['input_shapes']) - packed_npy_size = (18 + sum(policy_sizes)) * np.dtype(np.float32).itemsize - - def run_model(img_q, big_img_q, feat_q, desire_q, packed_npy_inputs): - packed_input = packed_npy_inputs.to(Device.DEFAULT) - Tensor.realize(packed_input) + def run_model(packed_npy_inputs, **state_inputs): + packed_input = packed_npy_inputs.to(Device.DEFAULT).realize() packed_npy_inputs = packed_input[:packed_npy_size].bitcast('float32') + inputs = {name: t.reshape(s) for (name, s), t in zip(shapes.items(), packed_npy_inputs.split(sizes), strict=True)} frame = packed_input[packed_npy_size:packed_npy_size + frame_copy_size] big_frame = packed_input[packed_npy_size + frame_copy_size:] - tfm, big_tfm, policy_inputs = packed_npy_inputs.split([9, 9, sum(policy_sizes)]) - warped = warp(tfm.reshape(3, 3), big_tfm.reshape(3, 3), frame, big_frame) - return run_policy(warped, img_q, big_img_q, feat_q, desire_q, policy_inputs) + inputs['new_img'] = warp(inputs.pop('tfm'), inputs.pop('big_tfm'), frame, big_frame) + inputs = {name: value.cast(input_shapes[name][1]) for name, value in inputs.items()} + outputs = {name: value.contiguous() for name, value in model_runner(inputs | state_inputs).items()} + Tensor.realize(*outputs.values()) + if state_pairs: + Tensor.realize(*(state_inputs[name].assign(outputs[next_name]) for name, next_name in state_pairs.items())) + return tuple(value for name, value in outputs.items() if name not in state_pairs.values()) return run_model -def compile_jit(jit, input_keys, make_queues, benchmark_runs): +def compile_jit(jit, make_queues, benchmark_runs): if benchmark_runs < 1: raise ValueError("benchmark_runs must be at least 1") @@ -245,7 +185,7 @@ def compile_jit(jit, input_keys, make_queues, benchmark_runs): v[:] = rng.integers(0, 256, size=v.shape, dtype=np.uint8) Device.default.synchronize() st = time.perf_counter() - outs = fn(**{k: input_queues[k] for k in input_keys}) + outs = fn(**input_queues) mt = time.perf_counter() Device.default.synchronize() et = time.perf_counter() @@ -300,7 +240,6 @@ if __name__ == "__main__": help='camera resolutions WxH (one or more)') p.add_argument('--onnx', required=True) p.add_argument('--output', required=True) - p.add_argument('--frame-skip', type=int, required=True) p.add_argument('--benchmark-runs', type=int, default=1, help='timed loaded-JIT runs for each correctness seed') args = p.parse_args() @@ -309,23 +248,24 @@ if __name__ == "__main__": model_w, model_h = args.model_size model_runner = OnnxRunner(model_path) + input_shapes = {name: (spec.shape, spec.dtype) for name, spec in model_runner.graph_inputs.items()} + state_pairs = {name: f'next_{name}' for name in input_shapes if f'next_{name}' in model_runner.graph_outputs} out = { 'metadata': make_metadata_dict(model_path), + 'input_shapes': input_shapes, + 'state_pairs': state_pairs, 'input_devices': {'model': Device.DEFAULT}, 'run_model': {}, } - run_policy = make_run_policy(model_runner, out['metadata'], args.frame_skip) - for cam_w, cam_h in args.camera_resolutions: nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) frame_copy_size = nv12_copy_size(nv12.stride, nv12.y_height, nv12.uv_height) - make_model_queues = partial(make_input_queues, out['metadata']['input_shapes'], args.frame_skip, + make_model_queues = partial(make_input_queues, input_shapes, state_pairs, frame_copy_size=frame_copy_size) warp = make_warp(nv12, model_w, model_h) - run_model_jit = TinyJit(make_run_model(warp, run_policy, out['metadata'], frame_copy_size), prune=True) - out['run_model'][(cam_w,cam_h)] = compile_jit(run_model_jit, MODELD_INPUTS, make_model_queues, - args.benchmark_runs) + run_model_jit = TinyJit(make_run_model(warp, model_runner, input_shapes, state_pairs, frame_copy_size), prune=True) + out['run_model'][(cam_w,cam_h)] = compile_jit(run_model_jit, make_model_queues, args.benchmark_runs) with open(args.output, "wb") as f: dump_oob(out, f) diff --git a/openpilot/selfdrive/modeld/modeld.py b/openpilot/selfdrive/modeld/modeld.py index a05361145a..994fda8b35 100755 --- a/openpilot/selfdrive/modeld/modeld.py +++ b/openpilot/selfdrive/modeld/modeld.py @@ -26,7 +26,7 @@ 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, should_stop, 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, nv12_copy_size, MODELD_INPUTS +from openpilot.selfdrive.modeld.compile_modeld import make_input_queues, nv12_copy_size from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_driving_model_data, fill_pose_msg, PublishState from openpilot.common.file_chunker import open_file_chunked from openpilot.selfdrive.modeld.constants import ModelConstants, Plan @@ -132,17 +132,17 @@ class ModelState: input_devices = jits['input_devices'] self.model_device = input_devices['model'] metadata = jits['metadata'] - self.input_shapes = metadata['input_shapes'] - self.vision_input_names = [k for k in self.input_shapes if 'img' in k] + self.input_shapes = jits['input_shapes'] + self.state_pairs = jits['state_pairs'] + self.vision_input_names = ('img', 'big_img') self.output_slices = metadata['output_slices'] self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32) self.chestnut = chestnut - self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ self.frame_copy_size = nv12_copy_size(*get_nv12_info(cam_w, cam_h)[:3]) self.input_queues, self.npy, self.frame_views = make_input_queues( - self.input_shapes, self.frame_skip, device=self.model_device, frame_copy_size=self.frame_copy_size) + self.input_shapes, self.state_pairs, device=self.model_device, frame_copy_size=self.frame_copy_size) self.parser = Parser() self.run_model = jits['run_model'][(cam_w,cam_h)] @@ -164,14 +164,13 @@ class ModelState: self.npy['tfm'][:,:] = transforms['img'][:,:] self.npy['big_tfm'][:,:] = transforms['big_img'][:,:] - outs, = self.run_model(**{k: self.input_queues[k] for k in MODELD_INPUTS}) + outs, = self.run_model(**self.input_queues) if after_enqueue is not None: after_enqueue() model_output = outs.numpy()[0] if self.chestnut and not np.all(np.isfinite(model_output)): raise RuntimeError("model output not finite") outputs_dict = self.parser.parse_outputs(self.slice_outputs(model_output, self.output_slices)) - self.npy['prev_feat'][:] = model_output[self.output_slices['hidden_state']] if SEND_RAW_PRED: outputs_dict['raw_pred'] = model_output.copy() @@ -183,7 +182,7 @@ class ModelState: dims = {'desire_pulse': ModelConstants.DESIRE_LEN, 'traffic_convention': 2, 'action_t': 2} self.run(dummy_frames, dict.fromkeys(self.vision_input_names, eye), {k: np.zeros(v, dtype=np.float32) for k, v in dims.items()}) self.input_queues, self.npy, self.frame_views = make_input_queues( - self.input_shapes, self.frame_skip, device=self.model_device, frame_copy_size=self.frame_copy_size) + self.input_shapes, self.state_pairs, device=self.model_device, frame_copy_size=self.frame_copy_size) self.prev_desire[:] = 0 diff --git a/openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx b/openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx index 13e3d32f47..3646adc744 100644 --- a/openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx +++ b/openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:09d080f36965bb2a0790500452bd328aa03c484d0222aa79d1ad9f021a522aec -size 766040736 +oid sha256:6fee5937923c74848df4a63f6239eb6331c6274dd4bdb7a5d6ec0388a8b543d5 +size 766018462 diff --git a/openpilot/selfdrive/modeld/models/driving_supercombo.onnx b/openpilot/selfdrive/modeld/models/driving_supercombo.onnx index f0672eab48..f030157ccd 100644 --- a/openpilot/selfdrive/modeld/models/driving_supercombo.onnx +++ b/openpilot/selfdrive/modeld/models/driving_supercombo.onnx @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:659727c4d4839adc4992a254409a54259a8756a743f2d567bf5fdc6579f8009b -size 60881999 +oid sha256:65a08adc31d5c456219687d99b7bf5e44d61dae2d49ea67850e76105c7248cce +size 60918562