get metadata at compile time

.gitignore

@@ -79,6 +79,7 @@ comma*.sh
 
 selfdrive/modeld/thneed/compile
 selfdrive/modeld/models/*.thneed
+selfdrive/modeld/models/*.pkl
 
 *.bz2
 

release/files_common

@@ -360,6 +360,7 @@ selfdrive/modeld/SConscript
 selfdrive/modeld/modeld.py
 selfdrive/modeld/parse_model_outputs.py
 selfdrive/modeld/fill_model_msg.py
+selfdrive/modeld/get_model_metadata.py
 selfdrive/modeld/navmodeld.py
 selfdrive/modeld/dmonitoringmodeld.py
 selfdrive/modeld/constants.py
@@ -373,7 +374,6 @@ selfdrive/modeld/models/commonmodel.cc
 selfdrive/modeld/models/commonmodel.h
 
 selfdrive/modeld/models/supercombo.onnx
-selfdrive/modeld/models/output_slices.pkl
 
 selfdrive/modeld/models/dmonitoring_model_q.dlc
 

selfdrive/modeld/SConscript

@@ -50,10 +50,14 @@ lenvCython.Program('runners/runmodel_pyx.so', 'runners/runmodel_pyx.pyx', LIBS=c
 lenvCython.Program('runners/snpemodel_pyx.so', 'runners/snpemodel_pyx.pyx', LIBS=[snpemodel_lib, snpe_lib, *cython_libs], FRAMEWORKS=frameworks, RPATH=snpe_rpath)
 lenvCython.Program('models/commonmodel_pyx.so', 'models/commonmodel_pyx.pyx', LIBS=[commonmodel_lib, *cython_libs], FRAMEWORKS=frameworks)
 
+# Get model metadata
+fn = File("models/supercombo").abspath
+cmd = f' python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
+files = sum([lenv.Glob("#"+x) for x in open(File("#release/files_common").abspath).read().split("\n") if x.endswith("get_model_metadata.py")], [])
+lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"]+files, cmd)
+
 # Build thneed model
 if arch == "larch64" or GetOption('pc_thneed'):
-  fn = File("models/supercombo").abspath
-
   tinygrad_opts = ["NATIVE_EXPLOG=1", "VALIDHACKS=1", "OPTLOCAL=1", "IMAGE=2", "GPU=1", "ENABLE_METHOD_CACHE=1"]
   if not GetOption('pc_thneed'):
     # use FLOAT16 on device for speed + don't cache the CL kernels for space

selfdrive/modeld/modeld.py

@@ -26,6 +26,8 @@ MODEL_PATHS = {
   ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed',
   ModelRunner.ONNX: Path(__file__).parent / 'models/supercombo.onnx'}
 
+METADATA_PATH = Path(__file__).parent / 'models/supercombo_metadata.pkl'
+
 class FrameMeta:
   frame_id: int = 0
   timestamp_sof: int = 0
@@ -55,10 +57,11 @@ class ModelState:
       'features_buffer': np.zeros(HISTORY_BUFFER_LEN * FEATURE_LEN, dtype=np.float32),
     }
 
-    with open(Path(__file__).parent / 'models/output_slices.pkl', 'rb') as f:
-      self.output_slices = pickle.load(f)
+    with open(METADATA_PATH, 'rb') as f:
+      model_metadata = pickle.load(f)
 
-    net_output_size = max([x.stop for x in self.output_slices.values()])
+    self.output_slices = model_metadata['output_slices']
+    net_output_size = model_metadata['output_shapes']['outputs'][1]
     self.output = np.zeros(net_output_size, dtype=np.float32)
 
     self.model = ModelRunner(MODEL_PATHS, self.output, Runtime.GPU, False, context)