You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
49 lines
1.8 KiB
49 lines
1.8 KiB
import sys, onnx
|
|
from tinygrad import Tensor, fetch, GlobalCounters
|
|
from tinygrad.uop import UOp
|
|
from tinygrad.frontend.onnx import OnnxRunner
|
|
from tinygrad.engine.grouper import get_kernelize_map
|
|
from tinygrad.engine.schedule import create_schedule_with_vars
|
|
from tinygrad.engine.realize import run_schedule
|
|
|
|
# NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
|
|
|
|
OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
|
|
OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
|
|
|
|
if __name__ == "__main__":
|
|
fn = fetch(OPENPILOT_MODEL)
|
|
onnx_file = fetch(OPENPILOT_MODEL)
|
|
onnx_model = onnx.load(onnx_file)
|
|
run_onnx = OnnxRunner(onnx_model)
|
|
|
|
inputs = run_onnx.get_empty_input_data("npy")
|
|
out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
|
|
root = out.lazydata
|
|
targets = [x.lazydata for x in inputs.values()]
|
|
print(targets)
|
|
|
|
# TODO: abstract this from gradient?
|
|
|
|
# compute the target path (top down)
|
|
in_target_path: dict[UOp, bool] = {}
|
|
for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
|
|
independent_set = {}
|
|
for u in root.toposort():
|
|
if in_target_path[u]:
|
|
for s in u.src:
|
|
if not in_target_path[s]:
|
|
independent_set[s] = None
|
|
independent = UOp.sink(*independent_set.keys())
|
|
kernelized = get_kernelize_map(independent)
|
|
independent = independent.substitute(kernelized)
|
|
schedule, var_vals, becomes_map = create_schedule_with_vars(independent)
|
|
run_schedule(schedule)
|
|
|
|
print("**** real ****")
|
|
GlobalCounters.reset()
|
|
out.lazydata = root.substitute(kernelized).substitute(becomes_map)
|
|
out.kernelize()
|
|
|
|
# realize
|
|
out.realize()
|
|
|