diff --git a/.gitignore b/.gitignore index f6d19d9e41..738a150b7e 100644 --- a/.gitignore +++ b/.gitignore @@ -44,7 +44,6 @@ bin/ config.json compile_commands.json compare_runtime*.html -selfdrive/modeld/models/tg_compiled_flags.json # build artifacts selfdrive/pandad/pandad diff --git a/selfdrive/modeld/SConscript b/selfdrive/modeld/SConscript index f45dd4e7c6..bad1cdd500 100644 --- a/selfdrive/modeld/SConscript +++ b/selfdrive/modeld/SConscript @@ -1,9 +1,6 @@ -import glob -import json import os -from SCons.Script import Value +import glob from openpilot.common.file_chunker import chunk_file, get_chunk_paths -from tinygrad import Device Import('env', 'arch') chunker_file = File("#common/file_chunker.py") @@ -16,51 +13,31 @@ tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + " def estimate_pickle_max_size(onnx_size): return 1.2 * onnx_size + 10 * 1024 * 1024 # 20% + 10MB is plenty +# compile warp # THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689 -# get fastest TG config -available = set(Device.get_available_devices()) -if 'CUDA' in available: - tg_backend = 'CUDA' - tg_flags = f'DEV={tg_backend}' -elif 'QCOM' in available: - tg_backend = 'QCOM' - tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0' -else: - tg_backend = 'CPU' - tg_flags = f'DEV={tg_backend} THREADS=0' - -def write_tg_compiled_flags(target, source, env): - with open(str(target[0]), "w") as f: - json.dump({"DEV": tg_backend}, f) - f.write("\n") - -compiled_flags_node = lenv.Command( - File("models/tg_compiled_flags.json").abspath, - tinygrad_files + [Value(tg_backend)], - write_tg_compiled_flags, -) - -# tinygrad calls brew which needs a $HOME in the env -mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else '' +tg_flags = { + 'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0', + 'Darwin': f'DEV=CPU THREADS=0 HOME={os.path.expanduser("~")}', # tinygrad calls brew which needs a $HOME in the env +}.get(arch, 'DEV=CPU:LLVM THREADS=0') # Get model metadata for model_name in ['driving_vision', 'driving_off_policy', 'driving_on_policy', 'dmonitoring_model']: fn = File(f"models/{model_name}").abspath script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)] - cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' - lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files + [compiled_flags_node], cmd) + cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx' + lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files, cmd) image_flag = { 'larch64': 'IMAGE=2', }.get(arch, 'IMAGE=0') script_files = [File(Dir("#selfdrive/modeld").File("compile_warp.py").abspath)] -compile_warp_cmd = f'{tg_flags} {mac_brew_string} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py ' +compile_warp_cmd = f'{tg_flags} python3 {Dir("#selfdrive/modeld").abspath}/compile_warp.py ' from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye warp_targets = [] for cam in [_ar_ox_fisheye, _os_fisheye]: w, h = cam.width, cam.height warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath] -lenv.Command(warp_targets, tinygrad_files + script_files + [compiled_flags_node], compile_warp_cmd) +lenv.Command(warp_targets, tinygrad_files + script_files, compile_warp_cmd) def tg_compile(flags, model_name): pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"' @@ -70,7 +47,7 @@ def tg_compile(flags, model_name): chunk_targets = get_chunk_paths(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path))) compile_node = lenv.Command( pkl, - [onnx_path] + tinygrad_files + [chunker_file, compiled_flags_node], + [onnx_path] + tinygrad_files + [chunker_file], f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {pkl}', ) def do_chunk(target, source, env): diff --git a/selfdrive/modeld/dmonitoringmodeld.py b/selfdrive/modeld/dmonitoringmodeld.py index d5901e8943..efd8214b9f 100755 --- a/selfdrive/modeld/dmonitoringmodeld.py +++ b/selfdrive/modeld/dmonitoringmodeld.py @@ -1,12 +1,12 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags -set_tinygrad_backend_from_compiled_flags() - +from openpilot.system.hardware import TICI +os.environ['DEV'] = 'QCOM' if TICI else 'CPU' from tinygrad.tensor import Tensor import time import pickle import numpy as np +from pathlib import Path from cereal import messaging from cereal.messaging import PubMaster, SubMaster @@ -21,8 +21,9 @@ from openpilot.selfdrive.modeld.parse_model_outputs import sigmoid, safe_exp PROCESS_NAME = "selfdrive.modeld.dmonitoringmodeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') -MODEL_PKL_PATH = MODELS_DIR / 'dmonitoring_model_tinygrad.pkl' -METADATA_PATH = MODELS_DIR / 'dmonitoring_model_metadata.pkl' +MODEL_PKL_PATH = Path(__file__).parent / 'models/dmonitoring_model_tinygrad.pkl' +METADATA_PATH = Path(__file__).parent / 'models/dmonitoring_model_metadata.pkl' +MODELS_DIR = Path(__file__).parent / 'models' class ModelState: inputs: dict[str, np.ndarray] diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index a93bdd23ef..73425b9e5a 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 import os -from openpilot.selfdrive.modeld.tinygrad_helpers import MODELS_DIR, set_tinygrad_backend_from_compiled_flags -set_tinygrad_backend_from_compiled_flags() - +from openpilot.system.hardware import TICI +os.environ['DEV'] = 'QCOM' if TICI else 'CPU' USBGPU = "USBGPU" in os.environ if USBGPU: os.environ['DEV'] = 'AMD' @@ -13,6 +12,7 @@ import pickle import numpy as np import cereal.messaging as messaging from cereal import car, log +from pathlib import Path from cereal.messaging import PubMaster, SubMaster from msgq.visionipc import VisionIpcClient, VisionStreamType, VisionBuf from opendbc.car.car_helpers import get_demo_car_params @@ -37,6 +37,7 @@ from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase PROCESS_NAME = "selfdrive.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') +MODELS_DIR = Path(__file__).parent / 'models' VISION_PKL_PATH = MODELS_DIR / 'driving_vision_tinygrad.pkl' VISION_METADATA_PATH = MODELS_DIR / 'driving_vision_metadata.pkl' ON_POLICY_PKL_PATH = MODELS_DIR / 'driving_on_policy_tinygrad.pkl' diff --git a/selfdrive/modeld/tinygrad_helpers.py b/selfdrive/modeld/tinygrad_helpers.py deleted file mode 100644 index e833c20968..0000000000 --- a/selfdrive/modeld/tinygrad_helpers.py +++ /dev/null @@ -1,11 +0,0 @@ -import json -import os -from pathlib import Path - -MODELS_DIR = Path(__file__).parent / 'models' -COMPILED_FLAGS_PATH = MODELS_DIR / 'tg_compiled_flags.json' - - -def set_tinygrad_backend_from_compiled_flags() -> None: - with open(COMPILED_FLAGS_PATH) as f: - os.environ['DEV'] = str(json.load(f)['DEV'])