diff --git a/selfdrive/modeld/runners/onnx_runner.py b/selfdrive/modeld/runners/onnx_runner.py index 3b19550e8ff41e..8a66c44917c112 100644 --- a/selfdrive/modeld/runners/onnx_runner.py +++ b/selfdrive/modeld/runners/onnx_runner.py @@ -9,29 +9,38 @@ import onnxruntime as ort # pylint: disable=import-error -def read(sz): +def read(sz, tf8=False): dd = [] gt = 0 - while gt < sz * 4: - st = os.read(0, sz * 4 - gt) + szof = 1 if tf8 else 4 + while gt < sz * szof: + st = os.read(0, sz * szof - gt) assert(len(st) > 0) dd.append(st) gt += len(st) - return np.frombuffer(b''.join(dd), dtype=np.float32) + r = np.frombuffer(b''.join(dd), dtype=np.uint8 if tf8 else np.float32).astype(np.float32) + if tf8: + r = r / 255. + return r def write(d): os.write(1, d.tobytes()) -def run_loop(m): +def run_loop(m, tf8_input=False): ishapes = [[1]+ii.shape[1:] for ii in m.get_inputs()] keys = [x.name for x in m.get_inputs()] + + # run once to initialize CUDA provider + if "CUDAExecutionProvider" in m.get_providers(): + m.run(None, dict(zip(keys, [np.zeros(shp, dtype=np.float32) for shp in ishapes]))) + print("ready to run onnx model", keys, ishapes, file=sys.stderr) while 1: inputs = [] - for shp in ishapes: + for k, shp in zip(keys, ishapes): ts = np.product(shp) #print("reshaping %s with offset %d" % (str(shp), offset), file=sys.stderr) - inputs.append(read(ts).reshape(shp)) + inputs.append(read(ts, (k=='input_img' and tf8_input)).reshape(shp)) ret = m.run(None, dict(zip(keys, inputs))) #print(ret, file=sys.stderr) for r in ret: @@ -39,6 +48,7 @@ def run_loop(m): if __name__ == "__main__": + print(sys.argv, file=sys.stderr) print("Onnx available providers: ", ort.get_available_providers(), file=sys.stderr) options = ort.SessionOptions() options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL @@ -54,7 +64,10 @@ def run_loop(m): options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL provider = 'CPUExecutionProvider' - print("Onnx selected provider: ", [provider], file=sys.stderr) - ort_session = ort.InferenceSession(sys.argv[1], options, providers=[provider]) - print("Onnx using ", ort_session.get_providers(), file=sys.stderr) - run_loop(ort_session) + try: + print("Onnx selected provider: ", [provider], file=sys.stderr) + ort_session = ort.InferenceSession(sys.argv[1], options, providers=[provider]) + print("Onnx using ", ort_session.get_providers(), file=sys.stderr) + run_loop(ort_session, tf8_input=("--use_tf8" in sys.argv)) + except KeyboardInterrupt: + pass \ No newline at end of file diff --git a/selfdrive/modeld/runners/onnxmodel.cc b/selfdrive/modeld/runners/onnxmodel.cc index cfcd2c954f83c3..c1185fe9937352 100644 --- a/selfdrive/modeld/runners/onnxmodel.cc +++ b/selfdrive/modeld/runners/onnxmodel.cc @@ -14,7 +14,7 @@ #include "common/swaglog.h" #include "common/util.h" -ONNXModel::ONNXModel(const char *path, float *_output, size_t _output_size, int runtime, bool _use_extra, bool _use_tf8) { +ONNXModel::ONNXModel(const char *path, float *_output, size_t _output_size, int runtime, bool _use_extra, bool _use_tf8, cl_context context) { LOGD("loading model %s", path); output = _output; diff --git a/selfdrive/modeld/runners/onnxmodel.h b/selfdrive/modeld/runners/onnxmodel.h index 8f01b0ce59e56d..2c9f6aec50699d 100644 --- a/selfdrive/modeld/runners/onnxmodel.h +++ b/selfdrive/modeld/runners/onnxmodel.h @@ -6,7 +6,7 @@ class ONNXModel : public RunModel { public: - ONNXModel(const char *path, float *output, size_t output_size, int runtime, bool use_extra = false, bool _use_tf8 = false); + ONNXModel(const char *path, float *output, size_t output_size, int runtime, bool use_extra = false, bool _use_tf8 = false, cl_context context = NULL); ~ONNXModel(); void addRecurrent(float *state, int state_size); void addDesire(float *state, int state_size);