Added ModelRunner helpers in modeld/runners/__init__.py (#29737)

* Added ModelRunner helper class

* import Runtime in __init__.py
pull/29742/head
Mitchell Goff 2 years ago committed by GitHub
parent 2dd7a6fc64
commit d5a55136a1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 16
      selfdrive/modeld/modeld.py
  2. 27
      selfdrive/modeld/runners/__init__.py

@ -1,5 +1,4 @@
#!/usr/bin/env python3
import os
import sys
import time
import numpy as np
@ -8,26 +7,21 @@ from typing import Dict, Optional
from setproctitle import setproctitle
from cereal.messaging import PubMaster, SubMaster
from cereal.visionipc import VisionIpcClient, VisionStreamType, VisionBuf
from openpilot.system.hardware import TICI
from openpilot.system.swaglog import cloudlog
from openpilot.common.params import Params
from openpilot.common.filter_simple import FirstOrderFilter
from openpilot.common.realtime import config_realtime_process
from openpilot.common.transformations.model import get_warp_matrix
from openpilot.selfdrive.modeld.runners.runmodel_pyx import Runtime
from openpilot.selfdrive.modeld.runners import ModelRunner, Runtime
from openpilot.selfdrive.modeld.models.commonmodel_pyx import ModelFrame, CLContext
from openpilot.selfdrive.modeld.models.driving_pyx import (
PublishState, create_model_msg, create_pose_msg,
FEATURE_LEN, HISTORY_BUFFER_LEN, DESIRE_LEN, TRAFFIC_CONVENTION_LEN, NAV_FEATURE_LEN, NAV_INSTRUCTION_LEN,
OUTPUT_SIZE, NET_OUTPUT_SIZE, MODEL_FREQ)
USE_THNEED = int(os.getenv('USE_THNEED', str(int(TICI))))
if USE_THNEED:
from selfdrive.modeld.runners.thneedmodel_pyx import ThneedModel as ModelRunner
else:
from selfdrive.modeld.runners.onnxmodel_pyx import ONNXModel as ModelRunner
MODEL_PATH = str(Path(__file__).parent / f"models/supercombo.{'thneed' if USE_THNEED else 'onnx'}")
MODEL_PATHS = {
ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed',
ModelRunner.ONNX: Path(__file__).parent / 'models/supercombo.onnx'}
class FrameMeta:
frame_id: int = 0
@ -59,7 +53,7 @@ class ModelState:
'feature_buffer': np.zeros(HISTORY_BUFFER_LEN * FEATURE_LEN, dtype=np.float32),
}
self.model = ModelRunner(MODEL_PATH, self.output, Runtime.GPU, False, context)
self.model = ModelRunner(MODEL_PATHS, self.output, Runtime.GPU, False, context)
self.model.addInput("input_imgs", None)
self.model.addInput("big_input_imgs", None)
for k,v in self.inputs.items():

@ -0,0 +1,27 @@
import os
from openpilot.system.hardware import TICI
from openpilot.selfdrive.modeld.runners.runmodel_pyx import RunModel, Runtime
assert Runtime
USE_THNEED = int(os.getenv('USE_THNEED', str(int(TICI))))
USE_SNPE = int(os.getenv('USE_SNPE', str(int(TICI))))
class ModelRunner(RunModel):
THNEED = 'THNEED'
SNPE = 'SNPE'
ONNX = 'ONNX'
def __new__(cls, paths, *args, **kwargs):
if ModelRunner.THNEED in paths and USE_THNEED:
from openpilot.selfdrive.modeld.runners.thneedmodel_pyx import ThneedModel as Runner
runner_type = ModelRunner.THNEED
elif ModelRunner.SNPE in paths and USE_SNPE:
from openpilot.selfdrive.modeld.runners.snpemodel_pyx import SNPEModel as Runner
runner_type = ModelRunner.SNPE
elif ModelRunner.ONNX in paths:
from openpilot.selfdrive.modeld.runners.onnxmodel_pyx import ONNXModel as Runner
runner_type = ModelRunner.ONNX
else:
raise Exception("Couldn't select a model runner, make sure to pass at least one valid model path")
return Runner(str(paths[runner_type]), *args, **kwargs)
Loading…
Cancel
Save