diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 88e72ea0b9..29ac9a47c9 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -1,13 +1,11 @@ import glob import json import os -from itertools import product from SCons.Script import Value from openpilot.common.file_chunker import chunk_file, get_chunk_paths 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 CompileConfig from tinygrad import Device from openpilot.system.hardware import HARDWARE, PC @@ -26,9 +24,6 @@ def get_camera_configs(): CAMERA_CONFIGS = get_camera_configs() -MODELD_CONFIGS = [CompileConfig(cam_w, cam_h, prepare_only, 'driving_') - for (cam_w, cam_h), prepare_only in product(CAMERA_CONFIGS, [True, False])] -DM_WARP_CONFIGS = [CompileConfig(cam_w, cam_h, True, 'dm_') for cam_w, cam_h in CAMERA_CONFIGS] chunker_file = File("#common/file_chunker.py") lenv = env.Clone() @@ -80,27 +75,28 @@ driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['dri model_w, model_h = MEDMODEL_INPUT_SIZE frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ -for cfg in MODELD_CONFIGS: - cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py ' - f'--model-size {model_w}x{model_h} ' - f'--nv12 {",".join(str(x) for x in cfg.nv12)} ' - f'--vision-onnx {File("models/driving_vision.onnx").abspath} ' - f'--policy-onnx {File("models/driving_policy.onnx").abspath} ' - f'--output {cfg.pkl_path} --frame-skip {frame_skip}' - + (' --prepare-only' if cfg.prepare_only else '')) - node = lenv.Command(cfg.pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + driving_metadata_deps + [chunker_file, compiled_flags_node], cmd) - onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps) - chunk_targets = get_chunk_paths(cfg.pkl_path, estimate_pickle_max_size(onnx_sizes_sum)) - def do_chunk(target, source, env, pkl=cfg.pkl_path, chunks=chunk_targets): - chunk_file(pkl, chunks) - lenv.Command(chunk_targets, node, do_chunk) +pkl_path = File("models/driving_tinygrad.pkl").abspath +camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS) +cmd = (f'{tg_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("models/driving_vision.onnx").abspath} ' + f'--policy-onnx {File("models/driving_policy.onnx").abspath} ' + f'--output {pkl_path} --frame-skip {frame_skip}') +node = lenv.Command(pkl_path, tinygrad_files + compile_modeld_script + driving_onnx_deps + driving_metadata_deps + [Value(camera_res_args), chunker_file, compiled_flags_node], cmd) +onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps) +chunk_targets = get_chunk_paths(pkl_path, estimate_pickle_max_size(onnx_sizes_sum)*2) # TODO make weight dedupe work on QCOM +def do_chunk(target, source, env, pkl=pkl_path, chunks=chunk_targets): + chunk_file(pkl, chunks) +lenv.Command(chunk_targets, node, do_chunk) dm_w, dm_h = DM_INPUT_SIZE -for cfg in DM_WARP_CONFIGS: +for cam_w, cam_h in CAMERA_CONFIGS: + dm_pkl_path = File(f"models/dm_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_dm_warp.py ' - f'--nv12 {",".join(str(x) for x in cfg.nv12)} --warp-to {dm_w}x{dm_h} ' - f'--output {cfg.pkl_path}') - lenv.Command(cfg.pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [compiled_flags_node], cmd) + f'--camera-resolution {cam_w}x{cam_h} --warp-to {dm_w}x{dm_h} ' + f'--output {dm_pkl_path}') + lenv.Command(dm_pkl_path, tinygrad_files + compile_dm_warp_script + compile_modeld_script + [compiled_flags_node], cmd) def tg_compile(flags, model_name): pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' diff --git a/selfdrive/modeld/compile_dm_warp.py b/selfdrive/modeld/compile_dm_warp.py index b035566535..548990ee15 100755 --- a/selfdrive/modeld/compile_dm_warp.py +++ b/selfdrive/modeld/compile_dm_warp.py @@ -7,7 +7,8 @@ from tinygrad.tensor import Tensor from tinygrad.device import Device from tinygrad.engine.jit import TinyJit -from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, warp_perspective_tinygrad, _parse_size, _parse_nv12 +from openpilot.system.camerad.cameras.nv12_info import get_nv12_info +from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, warp_perspective_tinygrad, _parse_size def make_warp_dm(nv12: NV12Frame, dm_w, dm_h): @@ -44,11 +45,12 @@ def compile_dm_warp(nv12: NV12Frame, dm_w, dm_h, pkl_path): if __name__ == "__main__": p = argparse.ArgumentParser() - p.add_argument('--nv12', type=_parse_nv12, required=True, - help=f'NV12 frame layout: {",".join(NV12Frame._fields)}') + p.add_argument('--camera-resolution', type=_parse_size, required=True, help='camera resolution WxH') p.add_argument('--warp-to', type=_parse_size, required=True, help='DM input WxH') p.add_argument('--output', required=True) args = p.parse_args() + cam_w, cam_h = args.camera_resolution + nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) dm_w, dm_h = args.warp_to - compile_dm_warp(args.nv12, dm_w, dm_h, args.output) + compile_dm_warp(nv12, dm_w, dm_h, args.output) diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index 2f97d8890c..f6ad398924 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 import argparse +import os import pickle import time from functools import partial @@ -10,7 +11,6 @@ from tinygrad.tensor import Tensor from tinygrad.helpers import Context from tinygrad.device import Device from tinygrad.engine.jit import TinyJit -from tinygrad.nn.onnx import OnnxRunner NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size']) @@ -158,21 +158,10 @@ def make_run_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, mode def compile_modeld(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, - vision_onnx, policy_onnx, pkl_path): - from get_model_metadata import metadata_path_for - + vision_runner, policy_runner, vision_features_slice, + vision_input_shapes, policy_input_shapes): print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") - vision_runner = OnnxRunner(vision_onnx) - policy_runner = OnnxRunner(policy_onnx) - - with open(metadata_path_for(vision_onnx), 'rb') as f: - vision_metadata = pickle.load(f) - vision_features_slice = vision_metadata['output_slices']['hidden_state'] - vision_input_shapes = vision_metadata['input_shapes'] - with open(metadata_path_for(policy_onnx), 'rb') as f: - policy_input_shapes = pickle.load(f)['input_shapes'] - _run = make_run_policy(vision_runner, policy_runner, nv12, model_w, model_h, vision_features_slice, frame_skip, prepare_only) run_policy_jit = TinyJit(_run, prune=True) @@ -216,13 +205,10 @@ def compile_modeld(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip, run_policy_jit, test_val, test_buffers = random_inputs_run_fn(run_policy_jit, SEED) print('pickle round trip') - with open(pkl_path, "wb") as f: - pickle.dump(run_policy_jit, f) - print(f" Saved to {pkl_path}") - with open(pkl_path, "rb") as f: - run_policy_jit = pickle.load(f) + run_policy_jit = pickle.loads(pickle.dumps(run_policy_jit)) random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers, expect_match=True) random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False) + return run_policy_jit def _parse_size(s): @@ -230,25 +216,42 @@ def _parse_size(s): return int(w), int(h) -def _parse_nv12(s): - parts = s.split(',') - assert len(parts) == len(NV12Frame._fields), \ - f"--nv12 expects {','.join(NV12Frame._fields)} (got {s!r})" - return NV12Frame(*(int(x) for x in parts)) - - if __name__ == "__main__": + from tinygrad.nn.onnx import OnnxRunner + from openpilot.system.camerad.cameras.nv12_info import get_nv12_info p = argparse.ArgumentParser() p.add_argument('--model-size', type=_parse_size, required=True, help='model input WxH') - p.add_argument('--nv12', type=_parse_nv12, required=True, - help=f'NV12 frame layout: {",".join(NV12Frame._fields)}') + 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('--policy-onnx', required=True) p.add_argument('--output', required=True) - p.add_argument('--prepare-only', action='store_true') p.add_argument('--frame-skip', type=int, required=True) args = p.parse_args() model_w, model_h = args.model_size - compile_modeld(args.nv12, model_w, model_h, args.prepare_only, args.frame_skip, - args.vision_onnx, args.policy_onnx, args.output) + + # init runners once so weights are shared + from get_model_metadata import metadata_path_for + vision_runner = OnnxRunner(args.vision_onnx) + policy_runner = OnnxRunner(args.policy_onnx) + with open(metadata_path_for(args.vision_onnx), 'rb') as f: + vision_metadata = pickle.load(f) + vision_features_slice = vision_metadata['output_slices']['hidden_state'] + vision_input_shapes = vision_metadata['input_shapes'] + with open(metadata_path_for(args.policy_onnx), 'rb') as f: + policy_input_shapes = pickle.load(f)['input_shapes'] + + out = {} + for cam_w, cam_h in args.camera_resolutions: + nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h)) + out[(cam_w,cam_h)] = { + name: compile_modeld(nv12, model_w, model_h, prepare_only, args.frame_skip, + vision_runner, policy_runner, vision_features_slice, + vision_input_shapes, policy_input_shapes) + for name, prepare_only in [('warp_enqueue', True), ('run_policy', False)] + } + + with open(args.output, "wb") as f: + pickle.dump(out, f) + print(f"Saved combined JIT to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)") diff --git a/selfdrive/modeld/dmonitoringmodeld.py b/selfdrive/modeld/dmonitoringmodeld.py index 7b073a5f50..618d39ca67 100755 --- a/selfdrive/modeld/dmonitoringmodeld.py +++ b/selfdrive/modeld/dmonitoringmodeld.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.helpers import MODELS_DIR, CompileConfig, set_tinygrad_backend_from_compiled_flags +from openpilot.selfdrive.modeld.helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags set_tinygrad_backend_from_compiled_flags() from tinygrad.tensor import Tensor @@ -44,7 +44,7 @@ class ModelState: self.tensor_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} self._blob_cache : dict[int, Tensor] = {} self.model_run = pickle.loads(read_file_chunked(str(MODEL_PKL_PATH))) - with open(CompileConfig(cam_w, cam_h, prefix='dm_', prepare_only=True).pkl_path, "rb") as f: + with open(MODELS_DIR / f'dm_warp_{cam_w}x{cam_h}_tinygrad.pkl', "rb") as f: self.image_warp = pickle.load(f) def run(self, buf: VisionBuf, calib: np.ndarray, transform: np.ndarray) -> tuple[np.ndarray, float]: diff --git a/selfdrive/modeld/helpers.py b/selfdrive/modeld/helpers.py index e5d731f34b..9c1a0d4689 100644 --- a/selfdrive/modeld/helpers.py +++ b/selfdrive/modeld/helpers.py @@ -1,10 +1,7 @@ import json import os -from dataclasses import dataclass from pathlib import Path -from openpilot.system.camerad.cameras.nv12_info import get_nv12_info - MODELS_DIR = Path(__file__).resolve().parent / 'models' COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json' @@ -13,19 +10,3 @@ def set_tinygrad_backend_from_compiled_flags() -> None: if os.path.isfile(COMPILED_FLAGS_PATH): with open(COMPILED_FLAGS_PATH) as f: os.environ['DEV'] = str(json.load(f)['DEV']) - - -@dataclass -class CompileConfig: - cam_w: int - cam_h: int - prepare_only: bool - prefix: str - - @property - def pkl_path(self): - return str(MODELS_DIR / f'{self.prefix}{"warp_" if self.prepare_only else ""}{self.cam_w}x{self.cam_h}_tinygrad.pkl') - - @property - def nv12(self): - return (self.cam_w, self.cam_h, *get_nv12_info(self.cam_w, self.cam_h)) diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 73ed19ec94..5784099a38 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.helpers import MODELS_DIR, CompileConfig, set_tinygrad_backend_from_compiled_flags +from openpilot.selfdrive.modeld.helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags set_tinygrad_backend_from_compiled_flags() USBGPU = "USBGPU" in os.environ @@ -100,8 +100,9 @@ class ModelState: self._blob_cache : dict[int, Tensor] = {} self.parser = Parser() self.frame_buf_params = {k: get_nv12_info(cam_w, cam_h) for k in ('img', 'big_img')} - self.run_policy = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=False).pkl_path)) - self.warp_enqueue = pickle.loads(read_file_chunked(CompileConfig(cam_w, cam_h, prefix='driving_', prepare_only=True).pkl_path)) + jits = pickle.loads(read_file_chunked(MODELS_DIR / 'driving_tinygrad.pkl'))[(cam_w,cam_h)] + self.run_policy = jits['run_policy'] + self.warp_enqueue = jits['warp_enqueue'] self.warp_enqueue( **self.input_queues, frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(),