IQ.Pilot Release Commit @ d2ce8a8
This commit is contained in:
144
iqpilot/selfdrive/iqmodeld/tools/compile_egpu_model.py
Normal file
144
iqpilot/selfdrive/iqmodeld/tools/compile_egpu_model.py
Normal file
@@ -0,0 +1,144 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pickle
|
||||
import time
|
||||
|
||||
os.environ.setdefault("DEV", "USB+AMD:LLVM")
|
||||
os.environ.setdefault("FLOAT16", "1")
|
||||
os.environ.setdefault("JIT_BATCH_SIZE", "0")
|
||||
os.environ.setdefault("GMMU", "0")
|
||||
|
||||
import numpy as np
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_helpers import egpu_pkl_path, local_onnx, patch_tinygrad_fetch_fw
|
||||
from iqpilot.selfdrive.iqmodeld.egpu_model import EGPU_MODELS, get_egpu_model, resolve_egpu_model
|
||||
from iqpilot.selfdrive.iqmodeld.temporal_state import MODEL_INPUT_SPEC, spec_from_meta
|
||||
|
||||
INPUT_SPEC = dict(MODEL_INPUT_SPEC)
|
||||
|
||||
patch_tinygrad_fetch_fw()
|
||||
|
||||
SEED = 42
|
||||
|
||||
|
||||
def set_input_spec(meta: dict) -> None:
|
||||
spec = spec_from_meta(meta)
|
||||
if spec is not None:
|
||||
INPUT_SPEC.clear()
|
||||
INPUT_SPEC.update(spec)
|
||||
|
||||
|
||||
def make_run_model(model_runner):
|
||||
def run_model(**inputs):
|
||||
out = next(iter(model_runner({k: inputs[k] for k in INPUT_SPEC}).values())).cast("float32")
|
||||
return out.reshape(-1),
|
||||
return run_model
|
||||
|
||||
|
||||
def _random_inputs(seed: int):
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.tensor import Tensor
|
||||
rng = np.random.default_rng(seed)
|
||||
out = {}
|
||||
for name, (shape, dtype) in INPUT_SPEC.items():
|
||||
if dtype == "uint8":
|
||||
arr = rng.integers(0, 256, shape).astype(np.uint8)
|
||||
else:
|
||||
arr = rng.standard_normal(shape).astype(np.float32)
|
||||
out[name] = Tensor(arr, device=Device.DEFAULT).realize()
|
||||
return out
|
||||
|
||||
|
||||
def _run(fn, seed: int) -> np.ndarray:
|
||||
from tinygrad.device import Device
|
||||
st = time.perf_counter()
|
||||
outs = fn(**_random_inputs(seed))
|
||||
Device.default.synchronize()
|
||||
print(f" run(seed={seed}) {(time.perf_counter() - st) * 1e3:6.1f} ms")
|
||||
return outs[0].numpy().reshape(-1)
|
||||
|
||||
|
||||
def compile_model(meta: dict, onnx_path: str, out_path: str) -> str:
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
|
||||
if meta.get("split"):
|
||||
raise RuntimeError(f"model {meta['key']} is a split model; eGPU v1 compiles fused models only")
|
||||
|
||||
jit = TinyJit(make_run_model(OnnxRunner(onnx_path)), prune=True)
|
||||
|
||||
print("capture + replay")
|
||||
for _ in range(2):
|
||||
baseline = _run(jit, SEED)
|
||||
if baseline.shape[0] != meta["output_len"]:
|
||||
raise RuntimeError(f"model output length {baseline.shape[0]} != registry {meta['output_len']}")
|
||||
if not np.isfinite(baseline).all():
|
||||
raise RuntimeError("compiled model produced non-finite outputs")
|
||||
|
||||
print("pickle round trip")
|
||||
jit = pickle.loads(pickle.dumps(jit))
|
||||
if not np.array_equal(_run(jit, SEED), baseline):
|
||||
raise RuntimeError("outputs differ from baseline after pickle round trip")
|
||||
if np.array_equal(_run(jit, SEED + 1), baseline):
|
||||
raise RuntimeError("outputs insensitive to inputs after pickle round trip")
|
||||
|
||||
from tinygrad.tensor import Tensor
|
||||
zeros = {name: Tensor(np.zeros(shape, dtype=dtype), device=Device.DEFAULT).realize()
|
||||
for name, (shape, dtype) in INPUT_SPEC.items()}
|
||||
flat = jit(**zeros)[0].numpy().reshape(-1)
|
||||
from iqpilot.selfdrive.iqmodeld.parser import PhaseParser
|
||||
from iqpilot.selfdrive.iqmodeld.tools.compile_supercombo import _slice_outputs, _validate_pose_outputs
|
||||
_validate_pose_outputs(PhaseParser().parse_vision_outputs(_slice_outputs(flat, meta["output_slices"])))
|
||||
|
||||
bundle = {
|
||||
"run_model": jit,
|
||||
"model_key": meta["key"],
|
||||
"model_sha256": meta["sha256"],
|
||||
"output_len": int(meta["output_len"]),
|
||||
"frame_skip": int(meta["frame_skip"]),
|
||||
"input_spec": {name: (tuple(shape), dtype) for name, (shape, dtype) in INPUT_SPEC.items()},
|
||||
"input_device": Device.DEFAULT,
|
||||
}
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
tmp = out_path + ".part"
|
||||
with open(tmp, "wb") as f:
|
||||
pickle.dump(bundle, f, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
os.replace(tmp, out_path)
|
||||
return out_path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--model", default=None, help=f"registry key, one of {sorted(EGPU_MODELS)}")
|
||||
p.add_argument("--onnx", default=None)
|
||||
p.add_argument("--output", default=None)
|
||||
args = p.parse_args()
|
||||
|
||||
if args.model is not None:
|
||||
if args.model in EGPU_MODELS:
|
||||
meta = get_egpu_model(args.model)
|
||||
else:
|
||||
from iqpilot.common.params import Params
|
||||
meta = resolve_egpu_model(Params(), args.model)
|
||||
if meta is None:
|
||||
raise SystemExit(f"unknown model {args.model!r}: not a built-in ({sorted(EGPU_MODELS)}) and not in the synced catalog")
|
||||
else:
|
||||
meta = get_egpu_model()
|
||||
set_input_spec(meta)
|
||||
|
||||
onnx_path = args.onnx or local_onnx(meta)
|
||||
if onnx_path is None or not os.path.isfile(onnx_path):
|
||||
raise SystemExit(f"onnx not found for {meta['key']}; pass --onnx or let iqegpumodeld download it first")
|
||||
|
||||
out = compile_model(meta, onnx_path, args.output or egpu_pkl_path(meta))
|
||||
print(f"saved eGPU jit to {out} ({os.path.getsize(out) / 1e6:.2f} MB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
9
iqpilot/selfdrive/iqmodeld/tools/compile_emac_warp.py
Normal file
9
iqpilot/selfdrive/iqmodeld/tools/compile_emac_warp.py
Normal file
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
"""
|
||||
from iqpilot.selfdrive.iqmodeld.tools.compile_warp import MODEL_SIZE, compile_warp, main
|
||||
|
||||
__all__ = ["MODEL_SIZE", "compile_warp", "main"]
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
119
iqpilot/selfdrive/iqmodeld/tools/compile_warp.py
Normal file
119
iqpilot/selfdrive/iqmodeld/tools/compile_warp.py
Normal file
@@ -0,0 +1,119 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
||||
|
||||
Compile the backend-neutral warp-only artifact: NV12 camera frames + 3x3
|
||||
transforms -> (2, 6, model_h/2, model_w/2) uint8 warped tensor, on the device
|
||||
GPU (QCOM). maciqmodeld runs this locally
|
||||
and feed the output to their backend, so the big model's image pipeline is
|
||||
bit-identical to comma's fused pkl warp stage.
|
||||
|
||||
Run ON the device (needs the QCOM backend):
|
||||
cd /data/openpilot && DEV=QCOM WARP_DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 \
|
||||
python3 iqpilot/selfdrive/iqmodeld/tools/compile_warp.py \
|
||||
--camera-resolutions 1928x1208 --output /data/models/emac_warp.pkl
|
||||
The artifact is then split per-resolution into Paths.model_root().
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import os
|
||||
import pickle
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
|
||||
SELFTEST_SEED = 20260817
|
||||
|
||||
from iqpilot.selfdrive.iqmodeld.temporal_state import DEFAULT_FRAME_SKIP, MODEL_INPUT_SPEC
|
||||
from iqpilot.selfdrive.iqmodeld.tools.compile_supercombo import (
|
||||
NV12Frame, WARP_INPUTS, compile_jit, make_random_images, make_warp, make_warp_input_queues,
|
||||
)
|
||||
|
||||
MODEL_SIZE = (MODEL_INPUT_SPEC["img"][0][3] * 2, MODEL_INPUT_SPEC["img"][0][2] * 2) # (512, 256)
|
||||
|
||||
|
||||
def _parse_size(s: str) -> tuple[int, int]:
|
||||
w, h = s.lower().split("x")
|
||||
return int(w), int(h)
|
||||
|
||||
|
||||
def compile_warp(cam_w: int, cam_h: int, out_path: str | None = None,
|
||||
frame_skip: int = DEFAULT_FRAME_SKIP) -> str:
|
||||
"""Compile the warp-only QCOM JIT for one camera resolution and write the pkl.
|
||||
Returns the artifact path. Callable from the workers so a fresh device
|
||||
self-provisions the warp instead of erroring — needs the QCOM backend."""
|
||||
# the QCOM warp env must be set before tinygrad is imported here
|
||||
os.environ.setdefault("DEV", "QCOM")
|
||||
os.environ.setdefault("WARP_DEV", "QCOM")
|
||||
os.environ.setdefault("IMAGE", "1")
|
||||
os.environ.setdefault("FLOAT16", "1")
|
||||
os.environ.setdefault("NOLOCALS", "1")
|
||||
os.environ.setdefault("JIT_BATCH_SIZE", "0")
|
||||
from tinygrad.engine.jit import TinyJit
|
||||
from iqpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
from iqpilot.system.hardware.hw import Paths
|
||||
|
||||
model_w, model_h = MODEL_SIZE
|
||||
input_shapes = {name: shape for name, (shape, _) in MODEL_INPUT_SPEC.items()}
|
||||
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
||||
make_random_warp_inputs = partial(make_random_images, keys=["frame", "big_frame"],
|
||||
shape=nv12.size, device=os.getenv("WARP_DEV"))
|
||||
warp_jit = TinyJit(make_warp(nv12, model_w, model_h, frame_skip), prune=True)
|
||||
make_warp_queues = partial(make_warp_input_queues, input_shapes, frame_skip)
|
||||
compiled = compile_jit(warp_jit, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)
|
||||
|
||||
# historical artifact name: already-provisioned devices keep their warp
|
||||
out_path = out_path or os.path.join(Paths.model_root(), f"emac_warp_{cam_w}x{cam_h}_tinygrad.pkl")
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
tmp = out_path + ".part"
|
||||
bundle = {(cam_w, cam_h): compiled, "frame_skip": frame_skip, "model_size": MODEL_SIZE}
|
||||
bundle["selftest"] = selftest_digest(compiled, cam_w, cam_h, nv12.size)
|
||||
with open(tmp, "wb") as f:
|
||||
pickle.dump(bundle, f)
|
||||
os.replace(tmp, out_path) # atomic: a reader never sees a half-written pkl
|
||||
return out_path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--camera-resolutions", type=_parse_size, nargs="+", default=[(1928, 1208)])
|
||||
p.add_argument("--output", default=None)
|
||||
p.add_argument("--frame-skip", type=int, default=DEFAULT_FRAME_SKIP)
|
||||
args = p.parse_args()
|
||||
for cam_w, cam_h in args.camera_resolutions:
|
||||
out = compile_warp(cam_w, cam_h, args.output, frame_skip=args.frame_skip)
|
||||
print(f"saved warp JIT to {out} ({os.path.getsize(out) / 1e6:.2f} MB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
def selftest_inputs(cam_w: int, cam_h: int, nv12_size: int):
|
||||
"""A fixed synthetic frame pair and pair of matrices. Deterministic so the
|
||||
digest is reproducible on the device that compiled the artifact."""
|
||||
rng = np.random.default_rng(SELFTEST_SEED)
|
||||
frame = rng.integers(0, 256, nv12_size, dtype=np.uint8)
|
||||
big_frame = rng.integers(0, 256, nv12_size, dtype=np.uint8)
|
||||
tfm = np.array([[0.7, 0.02, 300.0], [0.01, 0.7, 240.0], [0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
big_tfm = np.array([[0.5, 0.01, 380.0], [0.02, 0.5, 300.0], [0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
return frame, big_frame, tfm, big_tfm
|
||||
|
||||
|
||||
def selftest_digest(compiled, cam_w: int, cam_h: int, nv12_size: int) -> str:
|
||||
"""Hash the warp's output for a fixed input.
|
||||
|
||||
A warp artifact pinned to one tinygrad can still unpickle under another and
|
||||
then compute silently wrong, which reaches the model as a garbage image and
|
||||
looks like a bad model rather than a stale artifact. A version string cannot
|
||||
see that; running it can."""
|
||||
from tinygrad.tensor import Tensor
|
||||
frame, big_frame, tfm, big_tfm = selftest_inputs(cam_w, cam_h, nv12_size)
|
||||
dev = os.getenv("WARP_DEV") or "QCOM"
|
||||
out = compiled(tfm=Tensor(tfm, device="NPY").realize(),
|
||||
big_tfm=Tensor(big_tfm, device="NPY").realize(),
|
||||
frame=Tensor(frame, device=dev).realize(),
|
||||
big_frame=Tensor(big_frame, device=dev).realize())
|
||||
return hashlib.sha256(out.numpy().astype(np.uint8).tobytes()).hexdigest()
|
||||
Reference in New Issue
Block a user