diff --git a/scripts/reporter.py b/scripts/reporter.py index 9821a47ccf..199f1fae58 100755 --- a/scripts/reporter.py +++ b/scripts/reporter.py @@ -25,7 +25,9 @@ class MetadataOnnxPBParser(OnnxPBParser): def get_checkpoint(f): model = MetadataOnnxPBParser(f).parse() metadata = {prop["key"]: prop["value"] for prop in model["metadata_props"]} - return metadata['model_checkpoint'].split('/')[0] + # "" or "...//"; combined models list vision then policy + parts = metadata['model_checkpoint'].split('/') + return parts[-2] if len(parts) > 1 else parts[0] if __name__ == "__main__": @@ -37,8 +39,8 @@ if __name__ == "__main__": master_path = MASTER_PATH + MODEL_PATH + fn if os.path.exists(master_path): master = get_checkpoint(master_path) - master_col = f"[{master}](https://reporter.comma.life/experiment/{master})" + master_col = f"[{master}](https://reporter.comma.life/{master})" else: master_col = "N/A (new model)" pr = get_checkpoint(BASEDIR + MODEL_PATH + fn) - print("|", fn, "|", master_col, "|", f"[{pr}](https://reporter.comma.life/experiment/{pr})", "|") + print("|", fn, "|", master_col, "|", f"[{pr}](https://reporter.comma.life/{pr})", "|") diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 80a7df4515..74adf0d932 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -86,15 +86,14 @@ frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ for usbgpu in [False, True] if USBGPU else [False]: target_pkl_path = File(modeld_pkl_path(usbgpu)).abspath - file_prefix, cmd_flags = ('big_', usbgpu_tg_flags) if usbgpu else ('', tg_flags) - driving_onnx_deps = [p for m in [f'{file_prefix}driving_vision', f'{file_prefix}driving_on_policy'] - for p in get_existing_chunks(File(f"models/{m}.onnx").abspath)] + # BIG_INTO_SMALL=1 builds the default target from the big model, e.g. to test it without a USB GPU + file_prefix, cmd_flags = ('big_', usbgpu_tg_flags) if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '', tg_flags) + driving_onnx_deps = get_existing_chunks(File(f"models/{file_prefix}driving_supercombo.onnx").abspath) camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS) cmd = (f'{cmd_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py ' f'--model-size {model_w}x{model_h} ' f'--camera-resolutions {camera_res_args} ' - f'--vision-onnx {File(f"models/{file_prefix}driving_vision.onnx").abspath} ' - f'--on-policy-onnx {File(f"models/{file_prefix}driving_on_policy.onnx").abspath} ' + f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} ' f'--output {target_pkl_path} --frame-skip {frame_skip}') 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)) diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index bc8caca7e6..09e8d06c7c 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -132,26 +132,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)} +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, 24, 512) + # 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], fb[2])} 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) +def make_input_queues(input_shapes, frame_skip, device): + input_queues, npy = make_warp_input_queues(input_shapes, frame_skip, device) - fb = policy_input_shapes['features_buffer'] # (1, 25, 512) - dp = policy_input_shapes['desire_pulse'] # (1, 25, 8) + fb = input_shapes['features_buffer'] # (1, 24, 512), past features only; the model appends the current frame's feature + dp = input_shapes['desire_pulse'] # (1, 25, 8) - shapes, sizes = get_policy_npy_shapes(policy_input_shapes) + shapes, sizes = get_policy_npy_shapes(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(), + 'feat_q': Tensor(np.zeros((frame_skip * fb[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(), 'packed_npy_inputs': Tensor(packed_npy_inputs, device='NPY').realize(), }) @@ -189,29 +191,27 @@ def make_warp(nv12, model_w, model_h, frame_skip): return warp_enqueue -def make_run_policy(model_runners, model_metadata, frame_skip): +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) - vision_features_slice = model_metadata['vision']['output_slices']['hidden_state'] - npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['on_policy']['input_shapes']) + npy_shapes, npy_sizes = get_policy_npy_shapes(model_metadata['input_shapes']) 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, 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) - 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) + 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, 'desire_pulse': desire_buf, 'traffic_convention': traffic_convention, 'action_t': action_t, } - on_policy_out = next(iter(model_runners['on_policy'](inputs).values())).cast('float32') - return Tensor.cat(vision_out, on_policy_out, dim=1), + out = next(iter(model_runner(inputs).values())).cast('float32') + return out, return run_policy @@ -281,26 +281,21 @@ if __name__ == "__main__": p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH') p.add_argument('--camera-resolutions', type=_parse_size, nargs='+', required=True, help='camera resolutions WxH (one or more)') - p.add_argument('--vision-onnx', required=True) - p.add_argument('--on-policy-onnx', required=True) + p.add_argument('--onnx', required=True) p.add_argument('--output', required=True) p.add_argument('--frame-skip', type=int, required=True) args = p.parse_args() - model_paths = { - 'vision': read_file_chunked_to_shm(args.vision_onnx), - 'on_policy': read_file_chunked_to_shm(args.on_policy_onnx), - } + model_path = read_file_chunked_to_shm(args.onnx) model_w, model_h = args.model_size - 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()}} + model_runner = OnnxRunner(model_path) + out = {'metadata': make_metadata_dict(model_path)} - run_policy_jit = TinyJit(make_run_policy(model_runners, out['metadata'], args.frame_skip), prune=True) + run_policy_jit = TinyJit(make_run_policy(model_runner, out['metadata'], args.frame_skip), prune=True) - 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']) + make_policy_queues = partial(make_input_queues, out['metadata']['input_shapes'], args.frame_skip) + make_random_model_inputs = partial(make_random_images, keys=['img', 'big_img'], shape=out['metadata']['input_shapes']['img']) out['run_policy'] = compile_jit(run_policy_jit, make_random_model_inputs, POLICY_INPUTS, make_policy_queues) @@ -308,7 +303,7 @@ if __name__ == "__main__": 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) - make_warp_queues = partial(make_warp_input_queues, out['metadata']['vision']['input_shapes'], args.frame_skip) + make_warp_queues = partial(make_warp_input_queues, out['metadata']['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: diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 7505ddf182..372e8704a2 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -79,20 +79,15 @@ class ModelState: input_devices = get_tg_input_devices(PROCESS_NAME, usbgpu) self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV'] jits = pickle.loads(read_file_chunked(modeld_pkl_path(usbgpu))) - vision_metadata = jits['metadata']['vision'] - 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] - - policy_metadata = jits['metadata']['on_policy'] - self.policy_input_shapes = policy_metadata['input_shapes'] - self.policy_output_slices = policy_metadata['output_slices'] + 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.output_slices = metadata['output_slices'] 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, device=self.QUEUE_DEV) + self.input_queues, self.npy = make_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV) self.full_frames: dict[str, Tensor] = {} self._blob_cache: dict[int, Tensor] = {} self.parser = Parser() @@ -132,14 +127,13 @@ class ModelState: 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 = np.split(outs.numpy()[0], [self.vision_output_len]) - 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(on_policy_output, self.policy_output_slices)) - combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict} + model_output = outs.numpy()[0] + 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: - combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), on_policy_output.copy()]) - return combined_outputs_dict + outputs_dict['raw_pred'] = model_output.copy() + return outputs_dict def main(demo=False): diff --git a/selfdrive/modeld/models/big_driving_on_policy.onnx b/selfdrive/modeld/models/big_driving_on_policy.onnx deleted file mode 100644 index f7b49c018a..0000000000 --- a/selfdrive/modeld/models/big_driving_on_policy.onnx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:565e53c38dcd64c50dd3fe4d5ee1530213aeefd66c3f6b67ea6a72a32612a6bf -size 14061419 diff --git a/selfdrive/modeld/models/big_driving_supercombo.onnx b/selfdrive/modeld/models/big_driving_supercombo.onnx new file mode 100644 index 0000000000..de84325912 --- /dev/null +++ b/selfdrive/modeld/models/big_driving_supercombo.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:471f372751d5931e939320c211e35b7255f8fa8015125b3b3fd48ef43020257e +size 195490097 diff --git a/selfdrive/modeld/models/big_driving_vision.onnx b/selfdrive/modeld/models/big_driving_vision.onnx deleted file mode 100644 index d14f1969e0..0000000000 --- a/selfdrive/modeld/models/big_driving_vision.onnx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1f0cab5033fe9e3bc5e174a2e790fa277f7d9fc44c65822d734064d2f899a9a0 -size 296203378 diff --git a/selfdrive/modeld/models/driving_on_policy.onnx b/selfdrive/modeld/models/driving_on_policy.onnx deleted file mode 100644 index 611ae9fe85..0000000000 --- a/selfdrive/modeld/models/driving_on_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_supercombo.onnx b/selfdrive/modeld/models/driving_supercombo.onnx new file mode 100644 index 0000000000..f0672eab48 --- /dev/null +++ b/selfdrive/modeld/models/driving_supercombo.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:659727c4d4839adc4992a254409a54259a8756a743f2d567bf5fdc6579f8009b +size 60881999 diff --git a/selfdrive/modeld/models/driving_vision.onnx b/selfdrive/modeld/models/driving_vision.onnx deleted file mode 100644 index 6c9fc4c84d..0000000000 --- a/selfdrive/modeld/models/driving_vision.onnx +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ee29ee5bce84d1ce23e9ff381280de9b4e4d96d2934cd751740354884e112c66 -size 46877473 diff --git a/tinygrad_repo b/tinygrad_repo index 556defa0f7..5039d954f2 160000 --- a/tinygrad_repo +++ b/tinygrad_repo @@ -1 +1 @@ -Subproject commit 556defa0f78ed10b6ce675a2fa15a1c2521b5e93 +Subproject commit 5039d954f27fd4d4ebeb76951e9933f17985c1a4