forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Release Commit @ 0798119
This commit is contained in:
IQ.Lvbs history cleanup
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import (
|
||||
CUSTOM_MODEL_PATH, NumpyDict, ShapeDict, SliceDict,
|
||||
)
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
|
||||
from openpilot.iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
|
||||
|
||||
def _tinygrad_imports():
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.device import Device
|
||||
return Tensor, Device
|
||||
|
||||
|
||||
WARP_DEV = os.getenv('WARP_DEV')
|
||||
|
||||
|
||||
class TinygradFusedRunner(ModelRunner):
|
||||
"""Runs a fused warp+vision+policy pkl. Bundle ships one `driving_fused_*` artifact."""
|
||||
|
||||
uses_opencl_warp: bool = False
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self._constants = SplitModelConstants
|
||||
self._parser = PhaseParser()
|
||||
|
||||
if len(self.models) != 1:
|
||||
raise ValueError(f"fused bundle must have exactly one artifact, got {list(self.models)}")
|
||||
self._model_data = next(iter(self.models.values()))
|
||||
|
||||
pkl_path = os.path.join(CUSTOM_MODEL_PATH, self._model_data.model.artifact.fileName)
|
||||
with open(pkl_path, 'rb') as f:
|
||||
self._fused: dict[Any, Any] = pickle.load(f)
|
||||
|
||||
self._vision_meta = self._fused['metadata']['vision']
|
||||
self._on_meta = self._fused['metadata']['on_policy']
|
||||
self._off_meta = self._fused['metadata']['off_policy']
|
||||
self._run_policy = self._fused['run_policy']
|
||||
self._warp_jits: dict[tuple[int, int], Any] = {k: v for k, v in self._fused.items() if isinstance(k, tuple)}
|
||||
if not self._warp_jits:
|
||||
raise ValueError("fused pkl has no warp JITs")
|
||||
|
||||
self._frame_skip: int = int(self._fused.get('frame_skip', 4))
|
||||
|
||||
self._queues: dict[str, Any] | None = None
|
||||
self._npy_buffers: dict[str, np.ndarray] | None = None
|
||||
self._cam_resolution: tuple[int, int] | None = None
|
||||
self._blob_cache: dict[tuple[str, int], Any] = {}
|
||||
|
||||
def _frame_tensor(self, key, buf):
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
arr = np.frombuffer(buf.data, dtype=np.uint8)
|
||||
ck = (key, arr.ctypes.data)
|
||||
t = self._blob_cache.get(ck)
|
||||
if t is None:
|
||||
t = Tensor.from_blob(arr.ctypes.data, (arr.size,), dtype='uint8', device=Device.DEFAULT)
|
||||
self._blob_cache[ck] = t
|
||||
return t
|
||||
|
||||
@property
|
||||
def vision_input_names(self) -> list[str]:
|
||||
return ['img', 'big_img']
|
||||
|
||||
@property
|
||||
def input_shapes(self) -> ShapeDict:
|
||||
return {**self._vision_meta['input_shapes'], **self._on_meta['input_shapes']}
|
||||
|
||||
@property
|
||||
def output_slices(self) -> SliceDict:
|
||||
merged: SliceDict = {}
|
||||
for src in (self._vision_meta['output_slices'], self._on_meta['output_slices'], self._off_meta['output_slices']):
|
||||
merged.update({k: v for k, v in src.items() if k != 'pad'})
|
||||
return merged
|
||||
|
||||
def prepare_inputs(self, imgs_cl, numpy_inputs, frames):
|
||||
raise RuntimeError("fused runner has no OpenCL path; use run_fused()")
|
||||
|
||||
def _ensure_queues(self, cam_w: int, cam_h: int) -> None:
|
||||
if self._queues is not None and self._cam_resolution == (cam_w, cam_h):
|
||||
return
|
||||
if (cam_w, cam_h) not in self._warp_jits:
|
||||
raise RuntimeError(f"no warp JIT for {cam_w}x{cam_h}; have {sorted(self._warp_jits)}")
|
||||
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
img_shape = self._vision_meta['input_shapes']['img']
|
||||
fb = self._on_meta['input_shapes']['features_buffer']
|
||||
dp = self._on_meta['input_shapes']['desire_pulse']
|
||||
n_frames = img_shape[1] // 6
|
||||
img_buf_shape = (self._frame_skip * (n_frames - 1) + 1, 6, img_shape[2], img_shape[3])
|
||||
|
||||
zeros_u8 = lambda shp: Tensor(np.zeros(shp, dtype=np.uint8), device=Device.DEFAULT).contiguous().realize()
|
||||
zeros_f32 = lambda shp: Tensor(np.zeros(shp, dtype=np.float32), device=Device.DEFAULT).contiguous().realize()
|
||||
|
||||
self._queues = {
|
||||
'img_q': zeros_u8(img_buf_shape),
|
||||
'big_img_q': zeros_u8(img_buf_shape),
|
||||
'feat_q': zeros_f32((self._frame_skip * (fb[1] - 1) + 1, fb[0], fb[2])),
|
||||
'desire_q': zeros_f32((self._frame_skip * dp[1], dp[0], dp[2])),
|
||||
}
|
||||
# shapes must match the captured run_policy JIT inputs
|
||||
on_shapes = self._on_meta['input_shapes']
|
||||
self._npy_buffers = {
|
||||
'desire': np.zeros(dp[2], dtype=np.float32),
|
||||
'traffic_convention': np.zeros(on_shapes['traffic_convention'], dtype=np.float32),
|
||||
'action_t': np.zeros(on_shapes['action_t'], dtype=np.float32),
|
||||
'tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
'big_tfm': np.zeros((3, 3), dtype=np.float32),
|
||||
}
|
||||
self._cam_resolution = (cam_w, cam_h)
|
||||
|
||||
def run_fused(self, bufs: dict, transforms: dict[str, np.ndarray], numpy_inputs: NumpyDict) -> NumpyDict:
|
||||
"""warp + vision + policy in one pass from raw NV12 bufs + transform matrices."""
|
||||
Tensor, Device = _tinygrad_imports()
|
||||
|
||||
main_buf = bufs['img']
|
||||
self._ensure_queues(main_buf.width, main_buf.height)
|
||||
assert self._queues is not None and self._npy_buffers is not None
|
||||
|
||||
desire_key = next((k for k in numpy_inputs if k.startswith('desire')), None)
|
||||
if desire_key is not None:
|
||||
self._npy_buffers['desire'][:] = numpy_inputs[desire_key]
|
||||
if 'traffic_convention' in numpy_inputs:
|
||||
self._npy_buffers['traffic_convention'][:] = numpy_inputs['traffic_convention']
|
||||
if 'action_t' in numpy_inputs:
|
||||
self._npy_buffers['action_t'][:] = numpy_inputs['action_t']
|
||||
self._npy_buffers['tfm'][:] = transforms['img']
|
||||
self._npy_buffers['big_tfm'][:] = transforms['big_img']
|
||||
|
||||
npy = lambda key: Tensor(self._npy_buffers[key], device='NPY')
|
||||
|
||||
# frames go on the compute device to match the captured warp JIT
|
||||
frame = self._frame_tensor('img', bufs['img'])
|
||||
big_frame = self._frame_tensor('big_img', bufs['big_img'])
|
||||
|
||||
warp_jit = self._warp_jits[self._cam_resolution]
|
||||
img, big_img = warp_jit(img_q=self._queues['img_q'], big_img_q=self._queues['big_img_q'],
|
||||
tfm=npy('tfm'), big_tfm=npy('big_tfm'), frame=frame, big_frame=big_frame)
|
||||
|
||||
vision_out_t, on_out_t, off_out_t = self._run_policy(
|
||||
img=img, big_img=big_img, feat_q=self._queues['feat_q'], desire_q=self._queues['desire_q'],
|
||||
desire=npy('desire'), traffic_convention=npy('traffic_convention'), action_t=npy('action_t'))
|
||||
|
||||
# parse each model's output on its own sliced dict; parsing a merged dict
|
||||
# would run parse_dynamic_outputs twice and double-parse plan/lead
|
||||
def _slice(tensor_out, meta) -> NumpyDict:
|
||||
flat = tensor_out.numpy().flatten()
|
||||
return {k: flat[np.newaxis, sl] for k, sl in meta['output_slices'].items() if k != 'pad'}
|
||||
|
||||
parsed: NumpyDict = {}
|
||||
parsed.update(self._parser.parse_vision_outputs(_slice(vision_out_t, self._vision_meta)))
|
||||
parsed.update(self._parser.parse_policy_outputs(_slice(off_out_t, self._off_meta)))
|
||||
parsed.update(self._parser.parse_policy_outputs(_slice(on_out_t, self._on_meta)))
|
||||
return parsed
|
||||
|
||||
def _run_model(self) -> NumpyDict:
|
||||
raise RuntimeError("fused path goes through run_fused(), not _run_model()")
|
||||
Reference in New Issue
Block a user