mirror of
https://gitlvb.teallvbs.xyz/IQ.Lvbs/IQ.Pilot.git
synced 2026-07-25 13:32:04 +08:00
IQ.Lvbs CI [bot]
418 lines
17 KiB
Python
418 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import atexit
|
|
import os
|
|
import pickle
|
|
import time
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
|
|
import numpy as np
|
|
|
|
|
|
def _patch_firmware_fetch() -> None:
|
|
import hashlib
|
|
import pathlib
|
|
|
|
import zstandard
|
|
from tinygrad import helpers
|
|
|
|
if not hasattr(helpers, "fetch_fw"):
|
|
return
|
|
|
|
original_fetch = helpers.fetch_fw
|
|
|
|
def fetch_fw(path, name, sha256):
|
|
archive_path = pathlib.Path(f"/lib/firmware/{path}/{name}.zst")
|
|
if archive_path.is_file():
|
|
blob = zstandard.ZstdDecompressor().stream_reader(archive_path.read_bytes()).read()
|
|
if hashlib.sha256(blob).hexdigest() == sha256:
|
|
return blob
|
|
return original_fetch(path, name, sha256)
|
|
|
|
helpers.fetch_fw = fetch_fw
|
|
|
|
|
|
_patch_firmware_fetch()
|
|
|
|
from tinygrad.device import Device
|
|
from tinygrad.engine.jit import TinyJit
|
|
from tinygrad.helpers import Context
|
|
from tinygrad.nn.onnx import OnnxRunner
|
|
from tinygrad.tensor import Tensor
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CameraGeometry:
|
|
width: int
|
|
height: int
|
|
stride: int
|
|
y_height: int
|
|
uv_height: int
|
|
size: int
|
|
|
|
|
|
WARP_DEVICE = os.getenv("WARP_DEV")
|
|
|
|
|
|
def _read_shared_copy(path: str) -> str:
|
|
from openpilot.common.file_chunker import read_file_chunked
|
|
from openpilot.system.hardware.hw import Paths
|
|
|
|
shm_path = os.path.join(Paths.shm_path(), os.path.basename(path))
|
|
atexit.register(lambda: os.path.exists(shm_path) and os.remove(shm_path))
|
|
with open(shm_path, "wb") as handle:
|
|
handle.write(read_file_chunked(path))
|
|
return shm_path
|
|
|
|
|
|
def _parse_size(text: str) -> tuple[int, int]:
|
|
width, height = text.lower().split("x")
|
|
return int(width), int(height)
|
|
|
|
|
|
def _rand_u8_inputs(keys: list[str], shape, device=None):
|
|
return {key: Tensor.randint(shape, low=0, high=256, dtype="uint8", device=device).realize() for key in keys}
|
|
|
|
|
|
def _phase_desire_key(policy_shapes: dict[str, tuple[int, ...]]) -> str:
|
|
for key in policy_shapes:
|
|
if key.startswith("desire"):
|
|
return key
|
|
raise KeyError("No desire-like key found in policy shapes")
|
|
|
|
|
|
def _phase_image_keys(vision_shapes: dict[str, tuple[int, ...]]) -> tuple[str, str]:
|
|
names = sorted(name for name in vision_shapes if "img" in name)
|
|
road_key = next((name for name in names if "big" not in name), None)
|
|
wide_key = next((name for name in names if "big" in name), None)
|
|
if road_key is None or wide_key is None:
|
|
raise ValueError(f"Unable to resolve road/wide image keys from {list(vision_shapes)}")
|
|
return road_key, wide_key
|
|
|
|
|
|
def _base_policy_keys(policy_shapes: dict[str, tuple[int, ...]]) -> set[str]:
|
|
return {
|
|
_phase_desire_key(policy_shapes),
|
|
"features_buffer",
|
|
"traffic_convention",
|
|
"action_t",
|
|
}
|
|
|
|
|
|
def _common_policy_shapes(role_shapes: dict[str, dict[str, tuple[int, ...]]]) -> dict[str, tuple[int, ...]]:
|
|
first_role = next(iter(role_shapes))
|
|
baseline = role_shapes[first_role]
|
|
for role_name, shape_map in role_shapes.items():
|
|
if shape_map != baseline:
|
|
raise ValueError(f"Policy input shapes differ for role {role_name}")
|
|
return baseline
|
|
|
|
|
|
def _phase_frame_skip(policy_shapes: dict[str, tuple[int, ...]]) -> int:
|
|
feature_shape = policy_shapes.get("features_buffer")
|
|
if feature_shape is None:
|
|
return 1
|
|
history_length = feature_shape[1]
|
|
return 1 if history_length >= 99 else 4
|
|
|
|
|
|
def _project_pixels(src_flat, inverse_matrix, dst_shape, src_shape, stride_pad, border_fill_val=None):
|
|
dst_w, dst_h = dst_shape
|
|
src_h, src_w = src_shape
|
|
|
|
x_coords = Tensor.arange(dst_w, device=WARP_DEVICE).reshape(1, dst_w).expand(dst_h, dst_w).reshape(-1)
|
|
y_coords = Tensor.arange(dst_h, device=WARP_DEVICE).reshape(dst_h, 1).expand(dst_h, dst_w).reshape(-1)
|
|
|
|
src_x = inverse_matrix[0, 0] * x_coords + inverse_matrix[0, 1] * y_coords + inverse_matrix[0, 2]
|
|
src_y = inverse_matrix[1, 0] * x_coords + inverse_matrix[1, 1] * y_coords + inverse_matrix[1, 2]
|
|
scale = inverse_matrix[2, 0] * x_coords + inverse_matrix[2, 1] * y_coords + inverse_matrix[2, 2]
|
|
|
|
src_x = src_x / scale
|
|
src_y = src_y / scale
|
|
|
|
rounded_x = Tensor.round(src_x)
|
|
rounded_y = Tensor.round(src_y)
|
|
gather_x = rounded_x.clip(0, src_w - 1).cast("int")
|
|
gather_y = rounded_y.clip(0, src_h - 1).cast("int")
|
|
gather_index = gather_y * (src_w + stride_pad) + gather_x
|
|
sampled = src_flat[gather_index]
|
|
|
|
if border_fill_val is None:
|
|
return sampled
|
|
|
|
inside = ((rounded_x >= 0) & (rounded_x <= src_w - 1) & (rounded_y >= 0) & (rounded_y <= src_h - 1)).cast(sampled.dtype)
|
|
return sampled * inside + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - inside)
|
|
|
|
|
|
def _pack_nv12_planes(stacked_frame):
|
|
y_height = (stacked_frame.shape[0] * 2) // 3
|
|
frame_width = stacked_frame.shape[1]
|
|
return Tensor.cat(
|
|
stacked_frame[0:y_height:2, 0::2],
|
|
stacked_frame[1:y_height:2, 0::2],
|
|
stacked_frame[0:y_height:2, 1::2],
|
|
stacked_frame[1:y_height:2, 1::2],
|
|
stacked_frame[y_height:y_height + y_height // 4].reshape((y_height // 2, frame_width // 2)),
|
|
stacked_frame[y_height + y_height // 4:y_height + y_height // 2].reshape((y_height // 2, frame_width // 2)),
|
|
dim=0,
|
|
).reshape((6, y_height // 2, frame_width // 2))
|
|
|
|
|
|
def _warp_program(camera: CameraGeometry, model_w: int, model_h: int):
|
|
uv_offset = camera.stride * camera.y_height
|
|
stride_pad = camera.stride - camera.width
|
|
|
|
def prepare_frame(nv12_blob, inverse_matrix):
|
|
uv_matrix = inverse_matrix * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], device=WARP_DEVICE)
|
|
uv_plane = nv12_blob[uv_offset:uv_offset + camera.uv_height * camera.stride].reshape(camera.uv_height, camera.stride)
|
|
with Context(SPLIT_REDUCEOP=0):
|
|
y_plane = _project_pixels(nv12_blob[:camera.height * camera.stride], inverse_matrix, (model_w, model_h), (camera.height, camera.width), stride_pad).realize()
|
|
u_plane = _project_pixels(uv_plane[:camera.height // 2, :camera.width:2].flatten(), uv_matrix, (model_w // 2, model_h // 2), (camera.height // 2, camera.width // 2), 0).realize()
|
|
v_plane = _project_pixels(uv_plane[:camera.height // 2, 1:camera.width:2].flatten(), uv_matrix, (model_w // 2, model_h // 2), (camera.height // 2, camera.width // 2), 0).realize()
|
|
return _pack_nv12_planes(y_plane.cat(u_plane).cat(v_plane).reshape((model_h * 3 // 2, model_w)))
|
|
|
|
return prepare_frame
|
|
|
|
|
|
def _sample_sparse(queue_tensor, frame_stride):
|
|
return queue_tensor[::frame_stride].contiguous().flatten(0, 1).unsqueeze(0)
|
|
|
|
|
|
def _sample_desire(queue_tensor, frame_stride):
|
|
return queue_tensor.reshape(-1, frame_stride, *queue_tensor.shape[1:]).max(1).flatten(0, 1).unsqueeze(0)
|
|
|
|
|
|
def _roll_queue(queue_tensor, incoming, sampler):
|
|
queue_tensor.assign(queue_tensor[1:].cat(incoming, dim=0).contiguous())
|
|
return sampler(queue_tensor)
|
|
|
|
|
|
def _vision_queue_buffers(vision_shapes: dict[str, tuple[int, ...]], frame_stride: int, device):
|
|
road_key, _ = _phase_image_keys(vision_shapes)
|
|
image_shape = vision_shapes[road_key]
|
|
frame_history = image_shape[1] // 6
|
|
queue_depth = frame_stride * (frame_history - 1) + 1
|
|
frame_queue_shape = (queue_depth, 6, image_shape[2], image_shape[3])
|
|
|
|
numpy_state = {
|
|
"tfm": np.zeros((3, 3), dtype=np.float32),
|
|
"big_tfm": np.zeros((3, 3), dtype=np.float32),
|
|
}
|
|
tensor_state = {
|
|
"img_q": Tensor(np.zeros(frame_queue_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
|
"big_img_q": Tensor(np.zeros(frame_queue_shape, dtype=np.uint8), device=device).contiguous().realize(),
|
|
**{name: Tensor(value, device="NPY").realize() for name, value in numpy_state.items()},
|
|
}
|
|
return tensor_state, numpy_state
|
|
|
|
|
|
def _policy_queue_buffers(vision_shapes: dict[str, tuple[int, ...]], policy_shapes: dict[str, tuple[int, ...]], frame_stride: int, device):
|
|
tensor_state, numpy_state = _vision_queue_buffers(vision_shapes, frame_stride, device)
|
|
desired_key = _phase_desire_key(policy_shapes)
|
|
feature_shape = policy_shapes["features_buffer"]
|
|
desired_shape = policy_shapes[desired_key]
|
|
traffic_shape = policy_shapes["traffic_convention"]
|
|
action_shape = policy_shapes.get("action_t", traffic_shape)
|
|
|
|
numpy_policy = {
|
|
"desire": np.zeros(desired_shape[2], dtype=np.float32),
|
|
"traffic_convention": np.zeros(traffic_shape, dtype=np.float32),
|
|
"action_t": np.zeros(action_shape, dtype=np.float32),
|
|
}
|
|
for key, shape in policy_shapes.items():
|
|
if key not in _base_policy_keys(policy_shapes):
|
|
numpy_policy[key] = np.zeros(shape, dtype=np.float32)
|
|
|
|
numpy_state.update(numpy_policy)
|
|
tensor_state.update({
|
|
"feat_q": Tensor(np.zeros((frame_stride * (feature_shape[1] - 1) + 1, feature_shape[0], feature_shape[2]), dtype=np.float32), device=device).contiguous().realize(),
|
|
"desire_q": Tensor(np.zeros((frame_stride * desired_shape[1], desired_shape[0], desired_shape[2]), dtype=np.float32), device=device).contiguous().realize(),
|
|
**{name: Tensor(value, device="NPY").realize() for name, value in numpy_policy.items()},
|
|
})
|
|
return tensor_state, numpy_state
|
|
|
|
|
|
def _stage_program(camera: CameraGeometry, model_w: int, model_h: int, frame_stride: int):
|
|
prepare_frame = _warp_program(camera, model_w, model_h)
|
|
sparse_sampler = partial(_sample_sparse, frame_stride=frame_stride)
|
|
|
|
def stage_inputs(img_q, big_img_q, tfm, big_tfm, frame, big_frame):
|
|
tfm = tfm.to(WARP_DEVICE)
|
|
big_tfm = big_tfm.to(WARP_DEVICE)
|
|
Tensor.realize(tfm, big_tfm)
|
|
staged_main = prepare_frame(frame, tfm).unsqueeze(0).to(Device.DEFAULT)
|
|
staged_wide = prepare_frame(big_frame, big_tfm).unsqueeze(0).to(Device.DEFAULT)
|
|
return (
|
|
_roll_queue(img_q, staged_main, sparse_sampler),
|
|
_roll_queue(big_img_q, staged_wide, sparse_sampler),
|
|
)
|
|
|
|
return stage_inputs
|
|
|
|
|
|
def _role_executor(model_runners: dict[str, OnnxRunner], meta_by_role: dict[str, dict], frame_stride: int):
|
|
desired_sampler = partial(_sample_desire, frame_stride=frame_stride)
|
|
sparse_sampler = partial(_sample_sparse, frame_stride=frame_stride)
|
|
vision_hidden_slice = meta_by_role["vision"]["output_slices"]["hidden_state"]
|
|
policy_roles = [name for name in meta_by_role if name != "vision"]
|
|
policy_shapes = _common_policy_shapes({name: meta_by_role[name]["input_shapes"] for name in policy_roles})
|
|
desired_key = _phase_desire_key(policy_shapes)
|
|
road_key, wide_key = _phase_image_keys(meta_by_role["vision"]["input_shapes"])
|
|
extra_keys = [key for key in policy_shapes if key not in _base_policy_keys(policy_shapes)]
|
|
|
|
def execute_bundle(img, big_img, feat_q, desire_q, desire, traffic_convention, action_t, **extra):
|
|
desired_tensor = desire.to(Device.DEFAULT)
|
|
traffic_tensor = traffic_convention.to(Device.DEFAULT)
|
|
action_tensor = action_t.to(Device.DEFAULT)
|
|
extra_tensors = {key: extra[key].to(Device.DEFAULT) for key in extra_keys if key in extra}
|
|
Tensor.realize(desired_tensor, traffic_tensor, action_tensor, *extra_tensors.values())
|
|
|
|
desire_buffer = _roll_queue(desire_q, desired_tensor.reshape(1, 1, -1), desired_sampler)
|
|
vision_output = next(iter(model_runners["vision"]({road_key: img, wide_key: big_img}).values())).cast("float32")
|
|
hidden_state = vision_output[:, vision_hidden_slice].reshape(1, -1).unsqueeze(0)
|
|
feature_buffer = _roll_queue(feat_q, hidden_state, sparse_sampler)
|
|
|
|
common_inputs = {
|
|
"features_buffer": feature_buffer,
|
|
desired_key: desire_buffer,
|
|
"traffic_convention": traffic_tensor,
|
|
"action_t": action_tensor,
|
|
**extra_tensors,
|
|
}
|
|
|
|
role_outputs = []
|
|
for role_name in policy_roles:
|
|
role_outputs.append(next(iter(model_runners[role_name](common_inputs).values())).cast("float32"))
|
|
return (vision_output, *role_outputs)
|
|
|
|
return execute_bundle
|
|
|
|
|
|
def _capture_and_freeze(jit_runner, random_inputs_factory, queue_keys, queue_factory):
|
|
seed_value = 42
|
|
|
|
def validate(fn, baseline_outputs=None, baseline_buffers=None, expect_match=True, replay_seed=seed_value):
|
|
queue_tensors, numpy_values = queue_factory(Device.DEFAULT)
|
|
np.random.seed(replay_seed)
|
|
Tensor.manual_seed(replay_seed)
|
|
|
|
replay_count = 1 if (baseline_outputs is not None or baseline_buffers is not None) else 3
|
|
for pass_index in range(replay_count):
|
|
for value in numpy_values.values():
|
|
value[:] = np.random.randn(*value.shape).astype(value.dtype)
|
|
Device.default.synchronize()
|
|
random_inputs = random_inputs_factory()
|
|
start_time = time.perf_counter()
|
|
outputs = fn(**{name: queue_tensors[name] for name in queue_keys}, **random_inputs)
|
|
enqueue_time = time.perf_counter()
|
|
Device.default.synchronize()
|
|
total_time = time.perf_counter()
|
|
print(f" [{pass_index + 1}/{replay_count}] enqueue {(enqueue_time - start_time) * 1e3:6.2f} ms -- total {(total_time - start_time) * 1e3:6.2f} ms")
|
|
|
|
if pass_index == 0:
|
|
output_snapshot = [np.copy(value.numpy()) for value in outputs]
|
|
buffer_snapshot = [np.copy(value.numpy().copy()) for value in queue_tensors.values()]
|
|
|
|
if baseline_outputs is not None:
|
|
matches = all(np.array_equal(current, reference) for current, reference in zip(output_snapshot, baseline_outputs, strict=True))
|
|
assert matches == expect_match, f"outputs {'differ from' if expect_match else 'match'} baseline"
|
|
if baseline_buffers is not None:
|
|
matches = all(np.array_equal(current, reference) for current, reference in zip(buffer_snapshot, baseline_buffers, strict=True))
|
|
assert matches == expect_match, f"buffers {'differ from' if expect_match else 'match'} baseline"
|
|
|
|
return output_snapshot, buffer_snapshot
|
|
|
|
print("capture + replay")
|
|
baseline_outputs, baseline_buffers = validate(jit_runner)
|
|
print("pickle round trip")
|
|
frozen = pickle.loads(pickle.dumps(jit_runner))
|
|
validate(frozen, baseline_outputs, baseline_buffers, expect_match=True)
|
|
validate(frozen, baseline_outputs, baseline_buffers, expect_match=False, replay_seed=seed_value + 1)
|
|
return frozen
|
|
|
|
|
|
def _arg_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model-size", type=_parse_size, required=True, help="model input WxH")
|
|
parser.add_argument("--camera-resolutions", type=_parse_size, nargs="+", required=True, help="camera resolutions WxH")
|
|
parser.add_argument("--vision-onnx", required=True)
|
|
parser.add_argument("--policy-onnx")
|
|
parser.add_argument("--off-policy-onnx")
|
|
parser.add_argument("--on-policy-onnx")
|
|
parser.add_argument("--output", required=True)
|
|
parser.add_argument("--frame-skip", type=int)
|
|
return parser
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
from openpilot.iqpilot.selfdrive.iqmodeld.metadata import build_metadata_record
|
|
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
|
|
|
args = _arg_parser().parse_args(argv)
|
|
model_w, model_h = args.model_size
|
|
|
|
policy_specs = [
|
|
("policy", args.policy_onnx),
|
|
("off_policy", args.off_policy_onnx),
|
|
("on_policy", args.on_policy_onnx),
|
|
]
|
|
active_policy_specs = [(role, path) for role, path in policy_specs if path]
|
|
if not active_policy_specs:
|
|
raise SystemExit("At least one policy ONNX must be provided")
|
|
|
|
model_paths = {"vision": _read_shared_copy(args.vision_onnx)}
|
|
for role_name, onnx_path in active_policy_specs:
|
|
model_paths[role_name] = _read_shared_copy(onnx_path)
|
|
|
|
model_runners = {role_name: OnnxRunner(path) for role_name, path in model_paths.items()}
|
|
meta_by_role = {role_name: build_metadata_record(path) for role_name, path in model_paths.items()}
|
|
|
|
shared_policy_shapes = _common_policy_shapes({
|
|
role_name: meta_by_role[role_name]["input_shapes"] for role_name, _ in active_policy_specs
|
|
})
|
|
frame_stride = args.frame_skip if args.frame_skip is not None else _phase_frame_skip(shared_policy_shapes)
|
|
|
|
package: dict[Any, Any] = {
|
|
"meta_by_role": meta_by_role,
|
|
"roles": [role_name for role_name, _ in active_policy_specs],
|
|
"frame_stride": frame_stride,
|
|
}
|
|
|
|
executor_jit = TinyJit(_role_executor(model_runners, meta_by_role, frame_stride), prune=True)
|
|
queue_factory = partial(_policy_queue_buffers, meta_by_role["vision"]["input_shapes"], shared_policy_shapes, frame_stride)
|
|
image_shape = meta_by_role["vision"]["input_shapes"][_phase_image_keys(meta_by_role["vision"]["input_shapes"])[0]]
|
|
package["execute_bundle"] = _capture_and_freeze(
|
|
executor_jit,
|
|
partial(_rand_u8_inputs, keys=["img", "big_img"], shape=image_shape),
|
|
["feat_q", "desire_q", "desire", "traffic_convention", "action_t", *[k for k in shared_policy_shapes if k not in _base_policy_keys(shared_policy_shapes)]],
|
|
queue_factory,
|
|
)
|
|
|
|
for camera_width, camera_height in args.camera_resolutions:
|
|
camera = CameraGeometry(camera_width, camera_height, *get_nv12_info(camera_width, camera_height))
|
|
stage_jit = TinyJit(_stage_program(camera, model_w, model_h, frame_stride), prune=True)
|
|
package[(camera_width, camera_height)] = {
|
|
"stage_inputs": _capture_and_freeze(
|
|
stage_jit,
|
|
partial(_rand_u8_inputs, keys=["frame", "big_frame"], shape=camera.size, device=WARP_DEVICE),
|
|
["img_q", "big_img_q", "tfm", "big_tfm"],
|
|
partial(_vision_queue_buffers, meta_by_role["vision"]["input_shapes"], frame_stride),
|
|
)
|
|
}
|
|
|
|
with open(args.output, "wb") as handle:
|
|
pickle.dump(package, handle)
|
|
print(f"Saved combined split runtime to {args.output} ({os.path.getsize(args.output) / 1e6:.2f} MB)")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|