modeld: standalone compile script (#37851)

* modeld: standalone compile script

* cleanup

* frame skip

* rm last op import

* dm warp

* no graph break

* +x compile_dm_warp.py

* don't import tg before setting device

* compile_modeld exports metadata

* update help

* namedtuple

* lint

* Revert "compile_modeld exports metadata"

This reverts commit 93c3c223567b4d4a074c9071d7f734c56f5aedcc.

* import

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 + '"'

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)

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)

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

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)
 

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))

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

selfdrive/modeld/tinygrad_helpers.py

@@ -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'])