diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index 86b62d2c1..053eb157b 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -16,7 +16,6 @@ tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + " def estimate_pickle_max_size(onnx_size): return 1.2 * onnx_size + 10 * 1024 * 1024 # 20% + 10MB is plenty -# THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689 # get fastest TG config available = set(Device.get_available_devices()) if 'CUDA' in available: @@ -24,9 +23,10 @@ if 'CUDA' in available: tg_flags = f'DEV={tg_backend}' elif 'QCOM' in available: tg_backend = 'QCOM' - tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0' + tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1' else: tg_backend = 'CPU' if arch == 'Darwin' else 'CPU:LLVM' + # THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689 tg_flags = f'DEV={tg_backend} THREADS=0' def write_tg_compiled_flags(target, source, env): @@ -53,14 +53,23 @@ 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_warp.py").abspath)] -compile_warp_cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py ' -from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye -warp_targets = [] -for cam in [_ar_ox_fisheye, _os_fisheye]: - w, h = cam.width, cam.height - warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath] -lenv.Command(warp_targets, tinygrad_files + script_files + [compiled_flags_node], compile_warp_cmd) +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 ' +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: + 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_file(pkl, chunks) + lenv.Command(chunk_targets, compile_node, do_chunk) def tg_compile(flags, model_name): pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' @@ -81,7 +90,4 @@ def tg_compile(flags, model_name): do_chunk, ) -# Compile small models -for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']: - tg_compile(tg_flags, model_name) - +tg_compile(tg_flags, 'dmonitoring_model') diff --git a/selfdrive/modeld/compile_modeld.py b/selfdrive/modeld/compile_modeld.py new file mode 100755 index 000000000..469110a85 --- /dev/null +++ b/selfdrive/modeld/compile_modeld.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +import time +import pickle +from dataclasses import dataclass +from itertools import product +from functools import partial + +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 + +# https://github.com/tinygrad/tinygrad/issues/15682 +from tinygrad.uop.ops import UOp, Ops +_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] + + +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) + + +def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad): + w_dst, h_dst = dst_shape + h_src, w_src = src_shape + + x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1) + y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1) + + # inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather) + src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2] + src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2] + src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2] + + src_x = src_x / src_w + src_y = src_y / src_w + + x_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int') + y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int') + idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped + + return src_flat[idx] + + +def frames_to_tensor(frames): + H = (frames.shape[0] * 2) // 3 + W = frames.shape[1] + in_img1 = Tensor.cat(frames[0:H:2, 0::2], + frames[1:H:2, 0::2], + frames[0:H:2, 1::2], + frames[1:H:2, 1::2], + frames[H:H+H//4].reshape((H//2, W//2)), + frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2)) + 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) + uv_offset = stride * y_height + stride_pad = stride - cam_w + + def frame_prepare_tinygrad(input_frame, M_inv): + # UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling + M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]]) + # deinterleave NV12 UV plane (UVUV... -> separate U, V) + uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride) + with Context(SPLIT_REDUCEOP=0): + y = warp_perspective_tinygrad(input_frame[:cam_h*stride], + M_inv, (model_w, model_h), + (cam_h, cam_w), stride_pad).realize() + u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(), + M_inv_uv, (model_w//2, model_h//2), + (cam_h//2, cam_w//2), 0).realize() + v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(), + M_inv_uv, (model_w//2, model_h//2), + (cam_h//2, cam_w//2), 0).realize() + yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w)) + tensor = frames_to_tensor(yuv) + return tensor + return frame_prepare_tinygrad + + +def make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip): + img = vision_input_shapes['img'] # (1, 12, 128, 256) + n_frames = img[1] // 6 + img_buf_shape = (frame_skip * (n_frames - 1) + 1, 6, img[2], img[3]) + + 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) + + npy = { + 'desire': np.zeros(dp[2], dtype=np.float32), + 'traffic_convention': np.zeros(tc, dtype=np.float32), + 'tfm': np.zeros((3, 3), dtype=np.float32), + 'big_tfm': np.zeros((3, 3), dtype=np.float32), + } + input_queues = { + 'img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(), + 'big_img_q': Tensor.zeros(img_buf_shape, dtype='uint8').contiguous().realize(), + 'feat_q': Tensor.zeros(frame_skip * (fb[1] - 1) + 1, fb[0], fb[2]).contiguous().realize(), + 'desire_q': Tensor.zeros(frame_skip * dp[1], dp[0], dp[2]).contiguous().realize(), + **{k: Tensor(v, device='NPY').realize() for k, v in npy.items()}, + } + return input_queues, npy + + +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_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) + 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, + 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) + sample_skip_fn = partial(sample_skip, frame_skip=frame_skip) + sample_desire_fn = partial(sample_desire, frame_skip=frame_skip) + + def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame): + img = shift_and_sample(img_q, frame_prepare(frame, tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn) + big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn) + + if prepare_only: + return img, big_img + + vision_out = next(iter(vision_runner({'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) + desire_buf = shift_and_sample(desire_q, desire.to(Device.DEFAULT).reshape(1, 1, -1), sample_desire_fn) + + inputs = {'features_buffer': feat_buf, 'desire_pulse': desire_buf, 'traffic_convention': traffic_convention.to(Device.DEFAULT)} + policy_out = next(iter(policy_runner(inputs).values())).cast('float32') + + return vision_out, policy_out + 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 + + _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) + print(f"Compiling combined policy JIT for {cam_w}x{cam_h}...") + + vision_runner = OnnxRunner(MODELS_DIR / 'driving_vision.onnx') + policy_runner = OnnxRunner(MODELS_DIR / 'driving_policy.onnx') + + with open(MODELS_DIR / 'driving_vision_metadata.pkl', '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: + 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_policy_jit = TinyJit(_run, prune=True) + + N_RUNS = 3 + SEED = 42 + + def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True): + input_queues, npy = make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip) + np.random.seed(seed) + + for i in range(N_RUNS): + frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize() + big_frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize() + for v in npy.values(): + v[:] = np.random.randn(*v.shape).astype(v.dtype) + Device.default.synchronize() + st = time.perf_counter() + outs = fn(**input_queues, frame=frame, big_frame=big_frame) + mt = time.perf_counter() + for o in outs: + # .realize() not needed once jitted, but needed for unjitted fn + o.realize() + Device.default.synchronize() + et = time.perf_counter() + print(f" [{i+1}/{N_RUNS}] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms") + + val = [np.copy(v.numpy()) for v in outs] + buffers = [np.copy(v.numpy().copy()) for v in input_queues.values()] + + if test_val is not None: + match = all(np.array_equal(a, b) for a, b in zip(val, test_val, strict=True)) + assert match == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline (seed={seed})" + if test_buffers is not None: + match = all(np.array_equal(a, b) for a, b in zip(buffers, test_buffers, strict=True)) + assert match == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline (seed={seed})" + return fn, val, buffers + + print('run unjitted') + _, test_val, test_buffers = random_inputs_run_fn(_run, seed=SEED) + print('capture + replay') + run_policy_jit, _, _ = random_inputs_run_fn(run_policy_jit, SEED, test_val, test_buffers) + + 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) + 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) + + +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(np.random.randint(0, 256, yuv_size, dtype=np.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 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) + + +if __name__ == "__main__": + run_and_save_pickle() diff --git a/selfdrive/modeld/compile_warp.py b/selfdrive/modeld/compile_warp.py deleted file mode 100755 index 47511f2a2..000000000 --- a/selfdrive/modeld/compile_warp.py +++ /dev/null @@ -1,198 +0,0 @@ -#!/usr/bin/env python3 -import time -import pickle -import numpy as np -from pathlib import Path -from tinygrad.tensor import Tensor -from tinygrad.helpers import Context -from tinygrad.device import Device -from tinygrad.engine.jit import TinyJit - -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 - -MODELS_DIR = Path(__file__).parent / 'models' - -CAMERA_CONFIGS = [ - (_ar_ox_fisheye.width, _ar_ox_fisheye.height), # tici: 1928x1208 - (_os_fisheye.width, _os_fisheye.height), # mici: 1344x760 -] - -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) - -IMG_BUFFER_SHAPE = (30, MEDMODEL_INPUT_SIZE[1] // 2, MEDMODEL_INPUT_SIZE[0] // 2) - - -def warp_pkl_path(w, h): - return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' - - -def dm_warp_pkl_path(w, h): - return MODELS_DIR / f'dm_warp_{w}x{h}_tinygrad.pkl' - - -def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad): - w_dst, h_dst = dst_shape - h_src, w_src = src_shape - - x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1) - y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1) - - # inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather) - src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2] - src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2] - src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2] - - src_x = src_x / src_w - src_y = src_y / src_w - - x_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int') - y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int') - idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped - - return src_flat[idx] - - -def frames_to_tensor(frames, model_w, model_h): - H = (frames.shape[0] * 2) // 3 - W = frames.shape[1] - in_img1 = Tensor.cat(frames[0:H:2, 0::2], - frames[1:H:2, 0::2], - frames[0:H:2, 1::2], - frames[1:H:2, 1::2], - frames[H:H+H//4].reshape((H//2, W//2)), - frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2)) - 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) - uv_offset = stride * y_height - stride_pad = stride - cam_w - - def frame_prepare_tinygrad(input_frame, M_inv): - # UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling - M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]]) - # deinterleave NV12 UV plane (UVUV... -> separate U, V) - uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride) - with Context(SPLIT_REDUCEOP=0): - y = warp_perspective_tinygrad(input_frame[:cam_h*stride], - M_inv, (model_w, model_h), - (cam_h, cam_w), stride_pad).realize() - u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(), - M_inv_uv, (model_w//2, model_h//2), - (cam_h//2, cam_w//2), 0).realize() - v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(), - M_inv_uv, (model_w//2, model_h//2), - (cam_h//2, cam_w//2), 0).realize() - yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w)) - tensor = frames_to_tensor(yuv, model_w, model_h) - return tensor - return frame_prepare_tinygrad - - -def make_update_img_input(frame_prepare, model_w, model_h): - def update_img_input_tinygrad(frame_buffer, frame, M_inv): - M_inv = M_inv.to(Device.DEFAULT) - new_img = frame_prepare(frame, M_inv) - frame_buffer.assign(frame_buffer[6:].cat(new_img, dim=0).contiguous()) - return Tensor.cat(frame_buffer[:6], frame_buffer[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2) - return update_img_input_tinygrad - - -def make_update_both_imgs(frame_prepare, model_w, model_h): - update_img = make_update_img_input(frame_prepare, model_w, model_h) - - def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv, - calib_big_img_buffer, new_big_img, M_inv_big): - calib_img_pair = update_img(calib_img_buffer, new_img, M_inv) - calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big) - return calib_img_pair, calib_big_img_pair - return update_both_imgs_tinygrad - - -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) - 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 compile_modeld_warp(cam_w, cam_h): - model_w, model_h = MEDMODEL_INPUT_SIZE - _, _, _, yuv_size = get_nv12_info(cam_w, cam_h) - - print(f"Compiling modeld warp for {cam_w}x{cam_h}...") - - frame_prepare = make_frame_prepare(cam_w, cam_h, model_w, model_h) - update_both_imgs = make_update_both_imgs(frame_prepare, model_w, model_h) - update_img_jit = TinyJit(update_both_imgs, prune=True) - - full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() - big_full_buffer = Tensor.zeros(IMG_BUFFER_SHAPE, dtype='uint8').contiguous().realize() - for i in range(10): - img_inputs = [full_buffer, - Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), - Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - big_img_inputs = [big_full_buffer, - Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(), - Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')] - inputs = img_inputs + big_img_inputs - Device.default.synchronize() - - st = time.perf_counter() - _ = update_img_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") - - pkl_path = warp_pkl_path(cam_w, cam_h) - with open(pkl_path, "wb") as f: - pickle.dump(update_img_jit, f) - print(f" Saved to {pkl_path}") - - jit = pickle.load(open(pkl_path, "rb")) - jit(*inputs) - - -def compile_dm_warp(cam_w, cam_h): - 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(np.random.randint(0, 256, yuv_size, dtype=np.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") - - pkl_path = dm_warp_pkl_path(cam_w, cam_h) - with open(pkl_path, "wb") as f: - pickle.dump(warp_dm_jit, f) - print(f" Saved to {pkl_path}") - - -def run_and_save_pickle(): - for cam_w, cam_h in CAMERA_CONFIGS: - compile_modeld_warp(cam_w, cam_h) - compile_dm_warp(cam_w, cam_h) - - -if __name__ == "__main__": - run_and_save_pickle() diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 0c50d8beb..7a8a44298 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -26,6 +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.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 @@ -34,17 +35,13 @@ from openpilot.selfdrive.modeld.constants import ModelConstants, Plan PROCESS_NAME = "selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') -VISION_PKL_PATH = MODELS_DIR / 'driving_vision_tinygrad.pkl' VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl' -POLICY_PKL_PATH = MODELS_DIR / 'driving_policy_tinygrad.pkl' POLICY_METADATA_PATH = MODELS_DIR / 'driving_policy_metadata.pkl' LAT_SMOOTH_SECONDS = 0.0 LONG_SMOOTH_SECONDS = 0.3 MIN_LAT_CONTROL_SPEED = 0.3 -IMG_QUEUE_SHAPE = (6*(ModelConstants.MODEL_RUN_FREQ//ModelConstants.MODEL_CONTEXT_FREQ + 1), 128, 256) -assert IMG_QUEUE_SHAPE[0] == 30 def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, @@ -79,106 +76,36 @@ class FrameMeta: if vipc is not None: self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof -class InputQueues: - def __init__ (self, model_fps, env_fps, n_frames_input): - assert env_fps % model_fps == 0 - assert env_fps >= model_fps - self.model_fps = model_fps - self.env_fps = env_fps - self.n_frames_input = n_frames_input - - self.dtypes = {} - self.shapes = {} - self.q = {} - - def update_dtypes_and_shapes(self, input_dtypes, input_shapes) -> None: - self.dtypes.update(input_dtypes) - if self.env_fps == self.model_fps: - self.shapes.update(input_shapes) - else: - for k in input_shapes: - shape = list(input_shapes[k]) - if 'img' in k: - n_channels = shape[1] // self.n_frames_input - shape[1] = (self.env_fps // self.model_fps + (self.n_frames_input - 1)) * n_channels - else: - shape[1] = (self.env_fps // self.model_fps) * shape[1] - self.shapes[k] = tuple(shape) - - def reset(self) -> None: - self.q = {k: np.zeros(self.shapes[k], dtype=self.dtypes[k]) for k in self.dtypes.keys()} - - def enqueue(self, inputs:dict[str, np.ndarray]) -> None: - for k in inputs.keys(): - if inputs[k].dtype != self.dtypes[k]: - raise ValueError(f'supplied input <{k}({inputs[k].dtype})> has wrong dtype, expected {self.dtypes[k]}') - input_shape = list(self.shapes[k]) - input_shape[1] = -1 - single_input = inputs[k].reshape(tuple(input_shape)) - sz = single_input.shape[1] - self.q[k][:,:-sz] = self.q[k][:,sz:] - self.q[k][:,-sz:] = single_input - - def get(self, *names) -> dict[str, np.ndarray]: - if self.env_fps == self.model_fps: - return {k: self.q[k] for k in names} - else: - out = {} - for k in names: - shape = self.shapes[k] - if 'img' in k: - n_channels = shape[1] // (self.env_fps // self.model_fps + (self.n_frames_input - 1)) - out[k] = np.concatenate([self.q[k][:, s:s+n_channels] for s in np.linspace(0, shape[1] - n_channels, self.n_frames_input, dtype=int)], axis=1) - elif 'pulse' in k: - # any pulse within interval counts - out[k] = self.q[k].reshape((shape[0], shape[1] * self.model_fps // self.env_fps, self.env_fps // self.model_fps, -1)).max(axis=2) - else: - idxs = np.arange(-1, -shape[1], -self.env_fps // self.model_fps)[::-1] - out[k] = self.q[k][:, idxs] - return out class ModelState: - inputs: dict[str, np.ndarray] - output: np.ndarray prev_desire: np.ndarray # for tracking the rising edge of the pulse - def __init__(self): + def __init__(self, cam_w: int, cam_h: int): with open(VISION_METADATA_PATH, 'rb') as f: vision_metadata = pickle.load(f) 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'] - vision_output_size = vision_metadata['output_shapes']['outputs'][1] with open(POLICY_METADATA_PATH, 'rb') as f: policy_metadata = pickle.load(f) self.policy_input_shapes = policy_metadata['input_shapes'] self.policy_output_slices = policy_metadata['output_slices'] - policy_output_size = policy_metadata['output_shapes']['outputs'][1] self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32) - # policy inputs - self.numpy_inputs = {k: np.zeros(self.policy_input_shapes[k], dtype=np.float32) for k in self.policy_input_shapes} - self.full_input_queues = InputQueues(ModelConstants.MODEL_CONTEXT_FREQ, ModelConstants.MODEL_RUN_FREQ, ModelConstants.N_FRAMES) - for k in ['desire_pulse', 'features_buffer']: - self.full_input_queues.update_dtypes_and_shapes({k: self.numpy_inputs[k].dtype}, {k: self.numpy_inputs[k].shape}) - self.full_input_queues.reset() - - self.img_queues = {'img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize(), - 'big_img': Tensor.zeros(IMG_QUEUE_SHAPE, dtype='uint8').contiguous().realize()} + 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) self.full_frames : dict[str, Tensor] = {} self._blob_cache : dict[int, Tensor] = {} - self.transforms_np = {k: np.zeros((3,3), dtype=np.float32) for k in self.img_queues} - self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()} - self.vision_output = np.zeros(vision_output_size, dtype=np.float32) - self.policy_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} - self.policy_output = np.zeros(policy_output_size, dtype=np.float32) self.parser = Parser() - self.frame_buf_params : dict[str, tuple[int, int, int, int]] = {} - self.update_imgs = None - self.vision_run = pickle.loads(read_file_chunked(str(VISION_PKL_PATH))) - self.policy_run = pickle.loads(read_file_chunked(str(POLICY_PKL_PATH))) + 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)) + self.warp_enqueue( + **self.input_queues, + frame=Tensor.zeros(self.frame_buf_params['img'][3], dtype='uint8').contiguous().realize(), + big_frame=Tensor.zeros(self.frame_buf_params['big_img'][3], dtype='uint8').contiguous().realize()) def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]: parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()} @@ -186,18 +113,6 @@ class ModelState: def run(self, bufs: dict[str, VisionBuf], transforms: dict[str, np.ndarray], inputs: dict[str, np.ndarray], prepare_only: bool) -> dict[str, np.ndarray] | None: - # Model decides when action is completed, so desire input is just a pulse triggered on rising edge - inputs['desire_pulse'][0] = 0 - new_desire = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) - self.prev_desire[:] = inputs['desire_pulse'] - if self.update_imgs is None: - for key in bufs.keys(): - w, h = bufs[key].width, bufs[key].height - self.frame_buf_params[key] = get_nv12_info(w, h) - warp_path = MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl' - with open(warp_path, "rb") as f: - self.update_imgs = pickle.load(f) - for key in bufs.keys(): ptr = bufs[key].data.ctypes.data yuv_size = self.frame_buf_params[key][3] @@ -206,30 +121,31 @@ class ModelState: if cache_key not in self._blob_cache: self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8') self.full_frames[key] = self._blob_cache[cache_key] - for key in bufs.keys(): - self.transforms_np[key][:,:] = transforms[key][:,:] - out = self.update_imgs(self.img_queues['img'], self.full_frames['img'], self.transforms['img'], - self.img_queues['big_img'], self.full_frames['big_img'], self.transforms['big_img']) - vision_inputs = {'img': out[0], 'big_img': out[1]} + # Model decides when action is completed, so desire input is just a pulse triggered on rising edge + inputs['desire_pulse'][0] = 0 + self.npy['desire'][:] = np.where(inputs['desire_pulse'] - self.prev_desire > .99, inputs['desire_pulse'], 0) + self.prev_desire[:] = inputs['desire_pulse'] + self.npy['traffic_convention'][:] = inputs['traffic_convention'] + self.npy['tfm'][:,:] = transforms['img'][:,:] + self.npy['big_tfm'][:,:] = transforms['big_img'][:,:] if prepare_only: + self.warp_enqueue(**self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img']) return None - self.vision_output = self.vision_run(**vision_inputs).contiguous().realize().uop.base.buffer.numpy().flatten() - vision_outputs_dict = self.parser.parse_vision_outputs(self.slice_outputs(self.vision_output, self.vision_output_slices)) + vision_output, policy_output = self.run_policy( + **self.input_queues, frame=self.full_frames['img'], big_frame=self.full_frames['big_img'] + ) - self.full_input_queues.enqueue({'features_buffer': vision_outputs_dict['hidden_state'], 'desire_pulse': new_desire}) - for k in ['desire_pulse', 'features_buffer']: - self.numpy_inputs[k][:] = self.full_input_queues.get(k)[k] - self.numpy_inputs['traffic_convention'][:] = inputs['traffic_convention'] - - self.policy_output = self.policy_run(**self.policy_inputs).contiguous().realize().uop.base.buffer.numpy().flatten() - policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(self.policy_output, self.policy_output_slices)) + vision_output = vision_output.numpy().flatten() + policy_output = policy_output.numpy().flatten() + 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(policy_output, self.policy_output_slices)) combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict} - if SEND_RAW_PRED: - combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy()]) + if SEND_RAW_PRED: + combined_outputs_dict['raw_pred'] = np.concatenate([vision_output.copy(), policy_output.copy()]) return combined_outputs_dict @@ -241,11 +157,6 @@ def main(demo=False): # also need to move the aux USB interrupts for good timings config_realtime_process(7, 54) - st = time.monotonic() - cloudlog.warning("loading model") - model = ModelState() - cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") - # visionipc clients while True: available_streams = VisionIpcClient.available_streams("camerad", block=False) @@ -269,6 +180,11 @@ def main(demo=False): if use_extra_client: cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})") + st = time.monotonic() + cloudlog.warning("loading model") + model = ModelState(vipc_client_main.width, vipc_client_main.height) + cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting") + # messaging pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry"]) sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"]) diff --git a/selfdrive/modeld/tinygrad_helpers.py b/selfdrive/modeld/tinygrad_helpers.py index 49a6ed616..9c1a0d468 100644 --- a/selfdrive/modeld/tinygrad_helpers.py +++ b/selfdrive/modeld/tinygrad_helpers.py @@ -2,7 +2,7 @@ import json import os from pathlib import Path -MODELS_DIR = Path(__file__).parent / 'models' +MODELS_DIR = Path(__file__).resolve().parent / 'models' COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json'