diff --git a/selfdrive/modeld/runners/onnxmodel.py b/selfdrive/modeld/runners/onnxmodel.py index 2a870392d5..b2827a8619 100644 --- a/selfdrive/modeld/runners/onnxmodel.py +++ b/selfdrive/modeld/runners/onnxmodel.py @@ -1,39 +1,12 @@ -import onnx -import itertools import os +import onnx import sys import numpy as np from typing import Any from openpilot.selfdrive.modeld.runners.runmodel_pyx import RunModel +from openpilot.selfdrive.modeld.runners.ort_helpers import convert_fp16_to_fp32, ORT_TYPES_TO_NP_TYPES -ORT_TYPES_TO_NP_TYPES = {'tensor(float16)': np.float16, 'tensor(float)': np.float32, 'tensor(uint8)': np.uint8} - -def attributeproto_fp16_to_fp32(attr): - float32_list = np.frombuffer(attr.raw_data, dtype=np.float16) - attr.data_type = 1 - attr.raw_data = float32_list.astype(np.float32).tobytes() - -def convert_fp16_to_fp32(onnx_path_or_bytes): - if isinstance(onnx_path_or_bytes, bytes): - model = onnx.load_from_string(onnx_path_or_bytes) - elif isinstance(onnx_path_or_bytes, str): - model = onnx.load(onnx_path_or_bytes) - - for i in model.graph.initializer: - if i.data_type == 10: - attributeproto_fp16_to_fp32(i) - for i in itertools.chain(model.graph.input, model.graph.output): - if i.type.tensor_type.elem_type == 10: - i.type.tensor_type.elem_type = 1 - for i in model.graph.node: - if i.op_type == 'Cast' and i.attribute[0].i == 10: - i.attribute[0].i = 1 - for a in i.attribute: - if hasattr(a, 't'): - if a.t.data_type == 10: - attributeproto_fp16_to_fp32(a.t) - return model.SerializeToString() def create_ort_session(path, fp16_to_fp32): os.environ["OMP_NUM_THREADS"] = "4" @@ -56,7 +29,7 @@ def create_ort_session(path, fp16_to_fp32): options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL provider = 'CPUExecutionProvider' - model_data = convert_fp16_to_fp32(path) if fp16_to_fp32 else path + model_data = convert_fp16_to_fp32(onnx.load(path)) if fp16_to_fp32 else path print("Onnx selected provider: ", [provider], file=sys.stderr) ort_session = ort.InferenceSession(model_data, options, providers=[provider]) print("Onnx using ", ort_session.get_providers(), file=sys.stderr) diff --git a/selfdrive/modeld/runners/ort_helpers.py b/selfdrive/modeld/runners/ort_helpers.py new file mode 100644 index 0000000000..26afb03562 --- /dev/null +++ b/selfdrive/modeld/runners/ort_helpers.py @@ -0,0 +1,36 @@ +import onnx +import onnxruntime as ort +import numpy as np +import itertools + +ORT_TYPES_TO_NP_TYPES = {'tensor(float16)': np.float16, 'tensor(float)': np.float32, 'tensor(uint8)': np.uint8} + +def attributeproto_fp16_to_fp32(attr): + float32_list = np.frombuffer(attr.raw_data, dtype=np.float16) + attr.data_type = 1 + attr.raw_data = float32_list.astype(np.float32).tobytes() + +def convert_fp16_to_fp32(model): + for i in model.graph.initializer: + if i.data_type == 10: + attributeproto_fp16_to_fp32(i) + for i in itertools.chain(model.graph.input, model.graph.output): + if i.type.tensor_type.elem_type == 10: + i.type.tensor_type.elem_type = 1 + for i in model.graph.node: + if i.op_type == 'Cast' and i.attribute[0].i == 10: + i.attribute[0].i = 1 + for a in i.attribute: + if hasattr(a, 't'): + if a.t.data_type == 10: + attributeproto_fp16_to_fp32(a.t) + return model.SerializeToString() + + +def make_onnx_cpu_runner(model_path): + options = ort.SessionOptions() + options.intra_op_num_threads = 4 + options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + model_data = convert_fp16_to_fp32(onnx.load(model_path)) + return ort.InferenceSession(model_data, options, providers=['CPUExecutionProvider'])