179 lines
7.7 KiB
Python
179 lines
7.7 KiB
Python
"""
|
|
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 iqpilot.selfdrive.iqmodeld.models.runners.model_runner import (
|
|
CUSTOM_MODEL_PATH, NumpyDict, ShapeDict, SliceDict,
|
|
)
|
|
from iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelRunner
|
|
from iqpilot.selfdrive.iqmodeld.models.split_model_constants import SplitModelConstants
|
|
from 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):
|
|
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])),
|
|
}
|
|
on_shapes = self._on_meta['input_shapes']
|
|
captured = self._run_policy.captured
|
|
jit_shapes = {
|
|
name: tuple(int(s) for s in view.shape)
|
|
for name, (view, _vars, _dtype, _device) in zip(captured.expected_names, captured.expected_input_info)
|
|
}
|
|
|
|
def policy_input_shape(name):
|
|
shape = on_shapes.get(name, jit_shapes.get(name))
|
|
if shape is None:
|
|
raise ValueError(f"fused pkl declares no shape for policy input {name}")
|
|
return shape
|
|
|
|
self._npy_buffers = {
|
|
'desire': np.zeros(dp[2], dtype=np.float32),
|
|
'traffic_convention': np.zeros(policy_input_shape('traffic_convention'), dtype=np.float32),
|
|
'tfm': np.zeros((3, 3), dtype=np.float32),
|
|
'big_tfm': np.zeros((3, 3), dtype=np.float32),
|
|
}
|
|
if 'action_t' in jit_shapes:
|
|
self._npy_buffers['action_t'] = np.zeros(policy_input_shape('action_t'), 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:
|
|
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 and 'action_t' in self._npy_buffers:
|
|
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')
|
|
|
|
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)
|
|
|
|
policy_inputs = dict(
|
|
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'))
|
|
if 'action_t' in self._npy_buffers:
|
|
policy_inputs['action_t'] = npy('action_t')
|
|
vision_out_t, on_out_t, off_out_t = self._run_policy(**policy_inputs)
|
|
|
|
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()")
|