145 lines
4.9 KiB
Python
145 lines
4.9 KiB
Python
"""
|
|
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()
|