diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 053eb157b..76fef75ba 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -1,10 +1,23 @@ 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 +CAMERA_CONFIGS = [ + (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 + (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 +] +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] + Import('env', 'arch') chunker_file = File("#common/file_chunker.py") lenv = env.Clone() @@ -53,23 +66,35 @@ for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: image_flag = { 'larch64': 'IMAGE=2', }.get(arch, 'IMAGE=0') -script_files = [File(Dir("#selfdrive/modeld").File("compile_modeld.py").abspath)] -compile_modeld_cmd = f'{tg_flags} {mac_brew_string} {image_flag} python3 {Dir("#selfdrive/modeld").abspath}/compile_modeld.py ' +modeld_dir = Dir("#selfdrive/modeld").abspath +compile_modeld_script = [File(f"{modeld_dir}/compile_modeld.py")] +compile_dm_warp_script = [File(f"{modeld_dir}/compile_dm_warp.py")] driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_policy']] driving_metadata_deps = [File(f"models/{m}_metadata.pkl").abspath for m in ['driving_vision', 'driving_policy']] -from openpilot.selfdrive.modeld.compile_modeld import MODELD_CONFIGS, DM_WARP_CONFIGS -policy_pkls = [File(cfg.pkl_path).abspath for cfg in MODELD_CONFIGS] -modeld_targets = policy_pkls + [File(cfg.pkl_path).abspath for cfg in DM_WARP_CONFIGS] -compile_node = lenv.Command(modeld_targets, tinygrad_files + script_files + driving_onnx_deps + driving_metadata_deps + [chunker_file, compiled_flags_node], compile_modeld_cmd) - -# chunk the combined policy pkls -for policy_pkl in policy_pkls: +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} {image_flag} 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(policy_pkl, estimate_pickle_max_size(onnx_sizes_sum)) - def do_chunk(target, source, env, pkl=policy_pkl, chunks=chunk_targets): + 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, compile_node, do_chunk) + lenv.Command(chunk_targets, node, do_chunk) + +dm_w, dm_h = DM_INPUT_SIZE +for cfg in DM_WARP_CONFIGS: + cmd = (f'{tg_flags} {mac_brew_string} {image_flag} 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) 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 new file mode 100755 index 000000000..2713cccf4 --- /dev/null +++ b/selfdrive/modeld/compile_dm_warp.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +import argparse +import pickle +import time + +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 + + +def make_warp_dm(nv12: NV12Frame, dm_w, dm_h): + cam_w, cam_h, stride, _, _, _ = nv12 + stride_pad = stride - cam_w + + def warp_dm(input_frame, M_inv): + M_inv = M_inv.to(Device.DEFAULT).realize() + return warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, + (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w) + return warp_dm + + +def compile_dm_warp(nv12: NV12Frame, dm_w, dm_h, pkl_path): + print(f"Compiling DM warp for {nv12.width}x{nv12.height} -> {dm_w}x{dm_h}...") + + warp_dm_jit = TinyJit(make_warp_dm(nv12, dm_w, dm_h), prune=True) + + for i in range(10): + frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() + M_inv = Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY') + Device.default.synchronize() + st = time.perf_counter() + warp_dm_jit(frame, M_inv).realize() + mt = time.perf_counter() + Device.default.synchronize() + et = time.perf_counter() + print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") + + with open(pkl_path, "wb") as f: + pickle.dump(warp_dm_jit, f) + print(f" Saved to {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('--warp-to', type=_parse_size, required=True, help='DM input WxH') + p.add_argument('--output', required=True) + args = p.parse_args() + + dm_w, dm_h = args.warp_to + compile_dm_warp(args.nv12, dm_w, dm_h, args.output) diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py index c11f899b7..1e05ae30a 100755 --- a/selfdrive/modeld/compile_modeld.py +++ b/selfdrive/modeld/compile_modeld.py @@ -1,20 +1,16 @@ #!/usr/bin/env python3 -import time +import argparse import pickle -from dataclasses import dataclass -from itertools import product +import time from functools import partial +from collections import namedtuple import numpy as np from tinygrad.tensor import Tensor from tinygrad.helpers import Context from tinygrad.device import Device from tinygrad.engine.jit import TinyJit - -from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR -from openpilot.system.camerad.cameras.nv12_info import get_nv12_info -from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE -from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye +from tinygrad.nn.onnx import OnnxRunner # https://github.com/tinygrad/tinygrad/issues/15682 from tinygrad.uop.ops import UOp, Ops @@ -22,24 +18,7 @@ _orig = UOp.__reduce__ UOp.__reduce__ = lambda self: (UOp.unique, ()) if self.op is Ops.UNIQUE else _orig(self) -@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') - -CAMERA_CONFIGS = [ - (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 - (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 -] -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] - +NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size']) 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) @@ -79,8 +58,8 @@ def frames_to_tensor(frames): return in_img1 -def make_frame_prepare(cam_w, cam_h, model_w, model_h): - stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h) +def make_frame_prepare(nv12: NV12Frame, model_w, model_h): + cam_w, cam_h, stride, y_height, uv_height, _ = nv12 uv_offset = stride * y_height stride_pad = stride - cam_w @@ -143,21 +122,9 @@ 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_dm(cam_w, cam_h, dm_w, dm_h): - stride, y_height, _, _ = get_nv12_info(cam_w, cam_h) - stride_pad = stride - cam_w - - def warp_dm(input_frame, M_inv): - M_inv = M_inv.to(Device.DEFAULT).realize() - result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w) - return result - return warp_dm - - -def make_run_policy(vision_runner, policy_runner, cam_w, cam_h, +def make_run_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, model_h, vision_features_slice, frame_skip, prepare_only=False): - model_w, model_h = MEDMODEL_INPUT_SIZE - frame_prepare = make_frame_prepare(cam_w, cam_h, model_w, model_h) + frame_prepare = make_frame_prepare(nv12, model_w, model_h) sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) @@ -187,27 +154,24 @@ def make_run_policy(vision_runner, policy_runner, cam_w, cam_h, return run_policy -def compile_modeld(cam_w, cam_h, prepare_only, pkl_path): - from tinygrad.nn.onnx import OnnxRunner - from openpilot.selfdrive.modeld.constants import ModelConstants +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 - _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) - print(f"Compiling combined policy JIT for {cam_w}x{cam_h}...") + print(f"Compiling combined policy JIT for {nv12.width}x{nv12.height} (prepare_only={prepare_only})...") - vision_runner = OnnxRunner(MODELS_DIR / 'driving_vision.onnx') - policy_runner = OnnxRunner(MODELS_DIR / 'driving_policy.onnx') + vision_runner = OnnxRunner(vision_onnx) + policy_runner = OnnxRunner(policy_onnx) - with open(MODELS_DIR / 'driving_vision_metadata.pkl', 'rb') as f: + 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(MODELS_DIR / 'driving_policy_metadata.pkl', 'rb') as f: + with open(metadata_path_for(policy_onnx), 'rb') as f: policy_input_shapes = pickle.load(f)['input_shapes'] - frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ - - _run = make_run_policy(vision_runner, policy_runner, - cam_w, cam_h, vision_features_slice, frame_skip, prepare_only) + _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) N_RUNS = 3 @@ -219,8 +183,8 @@ def compile_modeld(cam_w, cam_h, prepare_only, pkl_path): Tensor.manual_seed(seed) for i in range(N_RUNS): - frame = Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize() - big_frame = Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize() + frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() + big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize() for v in npy.values(): v[:] = np.random.randn(*v.shape).astype(v.dtype) Device.default.synchronize() @@ -260,37 +224,30 @@ def compile_modeld(cam_w, cam_h, prepare_only, pkl_path): random_inputs_run_fn(run_policy_jit, SEED+1, test_val, test_buffers, expect_match=False) -def compile_dm_warp(cam_w, cam_h, pkl_path): - dm_w, dm_h = DM_INPUT_SIZE - _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) - - print(f"Compiling DM warp for {cam_w}x{cam_h}...") - - warp_dm = make_warp_dm(cam_w, cam_h, dm_w, dm_h) - warp_dm_jit = TinyJit(warp_dm, prune=True) - - for i in range(10): - inputs = [Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize(), - Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - Device.default.synchronize() - st = time.perf_counter() - warp_dm_jit(*inputs) - mt = time.perf_counter() - Device.default.synchronize() - et = time.perf_counter() - print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") - - with open(pkl_path, "wb") as f: - pickle.dump(warp_dm_jit, f) - print(f" Saved to {pkl_path}") +def _parse_size(s): + w, h = s.lower().split('x') + return int(w), int(h) -def run_and_save_pickle(): - for cfg in MODELD_CONFIGS: - compile_modeld(cfg.cam_w, cfg.cam_h, cfg.prepare_only, cfg.pkl_path) - for cfg in DM_WARP_CONFIGS: - compile_dm_warp(cfg.cam_w, cfg.cam_h, cfg.pkl_path) +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__": - run_and_save_pickle() + 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('--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) diff --git a/selfdrive/modeld/dmonitoringmodeld.py b/selfdrive/modeld/dmonitoringmodeld.py index d5901e894..9aaf6aa9e 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.tinygrad_helpers import MODELS_DIR, 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 diff --git a/selfdrive/modeld/get_model_metadata.py b/selfdrive/modeld/get_model_metadata.py index 838b1e9f4..ee08e9fb1 100755 --- a/selfdrive/modeld/get_model_metadata.py +++ b/selfdrive/modeld/get_model_metadata.py @@ -7,6 +7,10 @@ from typing import Any from tinygrad.nn.onnx import OnnxPBParser +def metadata_path_for(onnx_path) -> pathlib.Path: + p = pathlib.Path(onnx_path) + return p.parent / (p.stem + '_metadata.pkl') + class MetadataOnnxPBParser(OnnxPBParser): def _parse_ModelProto(self) -> dict: @@ -48,7 +52,7 @@ if __name__ == "__main__": 'output_shapes': dict(get_name_and_shape(x) for x in model["graph"]["output"]), } - metadata_path = model_path.parent / (model_path.stem + '_metadata.pkl') + metadata_path = metadata_path_for(model_path) with open(metadata_path, 'wb') as f: pickle.dump(metadata, f) diff --git a/selfdrive/modeld/helpers.py b/selfdrive/modeld/helpers.py new file mode 100644 index 000000000..e5d731f34 --- /dev/null +++ b/selfdrive/modeld/helpers.py @@ -0,0 +1,31 @@ +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' + + +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 7a8a44298..73ed19ec9 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.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags +from openpilot.selfdrive.modeld.helpers import MODELS_DIR, CompileConfig, set_tinygrad_backend_from_compiled_flags set_tinygrad_backend_from_compiled_flags() USBGPU = "USBGPU" in os.environ @@ -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, smooth_value, get_curvature_from_plan from openpilot.selfdrive.modeld.parse_model_outputs import Parser -from openpilot.selfdrive.modeld.compile_modeld import CompileConfig, make_input_queues +from openpilot.selfdrive.modeld.compile_modeld import make_input_queues from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState from openpilot.common.file_chunker import read_file_chunked from openpilot.selfdrive.modeld.constants import ModelConstants, Plan diff --git a/selfdrive/modeld/tinygrad_helpers.py b/selfdrive/modeld/tinygrad_helpers.py deleted file mode 100644 index 9c1a0d468..000000000 --- a/selfdrive/modeld/tinygrad_helpers.py +++ /dev/null @@ -1,12 +0,0 @@ -import json -import os -from pathlib import Path - -MODELS_DIR = Path(__file__).resolve().parent / 'models' -COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json' - - -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'])