mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-09-30 17:03:41 +08:00
Use tinygrad ONNX and warp compilers for driving and DM (#38922)
* modeld: use tinygrad ONNX and warp compilers for driving and DM * Bump tinygrad with current openpilot model compile tests * modeld: prioritize tinygrad compiler imports in build environment * Use compiler branch based on openpilot's pinned tinygrad * modeld: invoke tinygrad compilers without firmware patch
This commit is contained in:
Harald Schäfer
@@ -1,18 +1,18 @@
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from SCons.Script import Action, Value
|
||||
from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks
|
||||
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
|
||||
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
|
||||
from openpilot.selfdrive.modeld.helpers import TG_INPUT_DEVICES_PATH, chestnut_present, modeld_pkl_path
|
||||
from openpilot.selfdrive.modeld.helpers import chestnut_present, modeld_pkl_path
|
||||
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
||||
|
||||
|
||||
Import('env', 'arch')
|
||||
chunker_file = File("#openpilot/common/file_chunker.py")
|
||||
lenv = env.Clone()
|
||||
lenv.PrependENVPath('PYTHONPATH', Dir('#tinygrad_repo').abspath)
|
||||
|
||||
tinygrad_root = env.Dir("#").abspath
|
||||
tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
|
||||
@@ -26,62 +26,27 @@ def estimate_pickle_max_size(onnx_size):
|
||||
camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)]
|
||||
|
||||
if arch == 'comma_arm64':
|
||||
tg_backend = 'QCOM'
|
||||
tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
|
||||
tg_flags = 'DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
|
||||
else:
|
||||
tg_backend = 'METAL' if arch == 'Darwin' else 'CPU'
|
||||
# JIT=2 disables graph batching, which produces incorrect outputs after buffers change.
|
||||
tg_flags = 'DEV=METAL JIT=2' if arch == 'Darwin' else 'DEV=CPU:LLVM'
|
||||
|
||||
tg_devices = { # which device to put jit inputs to at runtime
|
||||
'openpilot.selfdrive.modeld.dmonitoringmodeld': {
|
||||
'default': {'DEV': tg_backend}
|
||||
},
|
||||
}
|
||||
|
||||
CHESTNUT = chestnut_present()
|
||||
if CHESTNUT:
|
||||
chestnut_tg_flags = 'DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32'
|
||||
# the USB+AMD GPU takes an exclusive flock; serialize all targets that touch it
|
||||
chestnut_lock = File("models/.chestnut.lock").abspath
|
||||
|
||||
def write_tg_devices(target, source, env):
|
||||
with open(str(target[0]), "w") as f:
|
||||
json.dump(tg_devices, f)
|
||||
f.write("\n")
|
||||
|
||||
tg_devices_node = lenv.Command(
|
||||
str(TG_INPUT_DEVICES_PATH),
|
||||
[Value(tg_devices)],
|
||||
write_tg_devices,
|
||||
)
|
||||
|
||||
# tinygrad calls brew which needs a $HOME in the env
|
||||
mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''
|
||||
|
||||
modeld_dir = Dir("#openpilot/selfdrive/modeld").abspath
|
||||
compile_modeld_script = [
|
||||
File(f"{modeld_dir}/compile_modeld.py"),
|
||||
File(f"{modeld_dir}/get_model_metadata.py"),
|
||||
File(f"{modeld_dir}/helpers.py"),
|
||||
File("#openpilot/common/hardware/hw.py"),
|
||||
]
|
||||
compile_warp_script = [File(f"{modeld_dir}/compile_warp.py"), File(f"{modeld_dir}/helpers.py"),
|
||||
File("#openpilot/system/camerad/cameras/nv12_info.py")]
|
||||
model_w, model_h = MEDMODEL_INPUT_SIZE
|
||||
warp_deps = [File("#openpilot/system/camerad/cameras/nv12_info.py")]
|
||||
compiler = 'python3 -m examples.openpilot'
|
||||
# CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it.
|
||||
taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else ''
|
||||
|
||||
for chestnut in [False, True] if CHESTNUT else [False]:
|
||||
target_pkl_path = File(modeld_pkl_path(chestnut)).abspath
|
||||
file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags)
|
||||
driving_onnx_deps = get_existing_chunks(File(f"models/{file_prefix}driving_supercombo.onnx").abspath)
|
||||
# CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it.
|
||||
taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else ''
|
||||
cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 {modeld_dir}/compile_modeld.py '
|
||||
f'--onnx {File(f"models/{file_prefix}driving_supercombo.onnx").abspath} '
|
||||
f'--output {target_pkl_path}')
|
||||
onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
|
||||
chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
|
||||
def do_compile(target, source, env, command=cmd, pkl=target_pkl_path, chunks=chunk_targets):
|
||||
def chestnut_action(command, pkl=None, chunks=()):
|
||||
def do_compile(target, source, env):
|
||||
from openpilot.system.hardware.chestnut.flash import link_up
|
||||
# chestnut can enumerate before its PCIe link is up due to varying 12V power behavior across cars
|
||||
for _ in range(10):
|
||||
@@ -95,58 +60,49 @@ for chestnut in [False, True] if CHESTNUT else [False]:
|
||||
return ret
|
||||
if chunks:
|
||||
chunk_file(pkl, chunks)
|
||||
return Action(do_compile, " [CHESTNUT] $TARGET")
|
||||
|
||||
def compile_model(onnx_path, pkl_path, flags, chestnut=False):
|
||||
onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath
|
||||
onnx_deps = get_existing_chunks(onnx_path)
|
||||
cmd = (f'{flags} {mac_brew_string} {taskset}{compiler}.compile_onnx '
|
||||
f'--onnx {onnx_path} --output {target_pkl_path}')
|
||||
onnx_sizes_sum = sum(os.path.getsize(f) for f in onnx_deps)
|
||||
chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
|
||||
def do_chunk(target, source, env, pkl=target_pkl_path, chunks=chunk_targets):
|
||||
chunk_file(pkl, chunks)
|
||||
actions = Action(do_compile, " [CHESTNUT] $TARGET") if chestnut else [cmd, Action(do_chunk, " [CHUNK] $TARGET")]
|
||||
actions = chestnut_action(cmd, target_pkl_path, chunk_targets) if chestnut else [cmd, Action(do_chunk, " [CHUNK] $TARGET")]
|
||||
node = lenv.Command(
|
||||
chunk_targets,
|
||||
tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file],
|
||||
tinygrad_files + onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file],
|
||||
actions,
|
||||
)
|
||||
if chestnut:
|
||||
lenv.SideEffect(chestnut_lock, node)
|
||||
|
||||
compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygrad.pkl', tg_flags)
|
||||
|
||||
model_w, model_h = MEDMODEL_INPUT_SIZE
|
||||
for chestnut in [False, True] if CHESTNUT else [False]:
|
||||
file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags)
|
||||
compile_model(f'models/{file_prefix}driving_supercombo.onnx', modeld_pkl_path(chestnut), cmd_flags, chestnut)
|
||||
for cam_w, cam_h in camera_configs:
|
||||
warp_pkl_path = File(f"models/{file_prefix}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
|
||||
stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
|
||||
cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 {modeld_dir}/compile_warp.py '
|
||||
cmd = (f'{cmd_flags} {mac_brew_string} {taskset}{compiler}.compile_warp '
|
||||
f'--camera-resolution {cam_w}x{cam_h} --warp-to {model_w}x{model_h} --layout yuv420 '
|
||||
f'--frames 2 --frame-size {stride * (y_height + uv_height)} '
|
||||
f'--frames 2 --stride {stride} --uv-offset {stride * y_height} --frame-size {stride * (y_height + uv_height)} '
|
||||
f'--output {warp_pkl_path}')
|
||||
def do_compile_warp(target, source, env, command=cmd):
|
||||
return do_compile(target, source, env, command=command, chunks=())
|
||||
action = Action(do_compile_warp, " [CHESTNUT] $TARGET") if chestnut else cmd
|
||||
node = lenv.Command(warp_pkl_path, tinygrad_files + compile_warp_script + [Value(cmd)], action)
|
||||
action = chestnut_action(cmd) if chestnut else cmd
|
||||
node = lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], action)
|
||||
if chestnut:
|
||||
lenv.SideEffect(chestnut_lock, node)
|
||||
|
||||
# get model metadata
|
||||
fn = File(f"models/dmonitoring_model").abspath
|
||||
script_files = [File(Dir("#openpilot/selfdrive/modeld").File("get_model_metadata.py").abspath)]
|
||||
cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#openpilot/selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
|
||||
lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [tg_devices_node], cmd)
|
||||
|
||||
dm_w, dm_h = DM_INPUT_SIZE
|
||||
for cam_w, cam_h in camera_configs:
|
||||
dm_pkl_path = File(f"models/dm_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
|
||||
cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_warp.py '
|
||||
stride, y_height, _, frame_size = get_nv12_info(cam_w, cam_h)
|
||||
cmd = (f'{tg_flags} {mac_brew_string} {compiler}.compile_warp '
|
||||
f'--camera-resolution {cam_w}x{cam_h} --warp-to {dm_w}x{dm_h} --layout luma --border-fill 16 --transform-device NPY '
|
||||
f'--output {dm_pkl_path}')
|
||||
lenv.Command(dm_pkl_path, tinygrad_files + compile_warp_script + [tg_devices_node], cmd)
|
||||
|
||||
def tg_compile(flags, model_name):
|
||||
pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
|
||||
fn = File(f"models/{model_name}").abspath
|
||||
pkl = fn + "_tinygrad.pkl"
|
||||
onnx_path = fn + ".onnx"
|
||||
chunk_targets = get_chunk_targets(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path)))
|
||||
def do_chunk(target, source, env):
|
||||
chunk_file(pkl, chunk_targets)
|
||||
return lenv.Command(
|
||||
chunk_targets,
|
||||
[onnx_path] + tinygrad_files + [Value(chunk_targets), chunker_file, tg_devices_node],
|
||||
[f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}',
|
||||
Action(do_chunk, " [CHUNK] $TARGET")],
|
||||
)
|
||||
|
||||
tg_compile(tg_flags, 'dmonitoring_model')
|
||||
f'--stride {stride} --uv-offset {stride * y_height} --frame-size {frame_size} --output {dm_pkl_path}')
|
||||
lenv.Command(dm_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd)
|
||||
|
||||
Reference in New Issue
Block a user