63 lines
1.8 KiB
Python
63 lines
1.8 KiB
Python
"""
|
|
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
|
|
"""
|
|
import itertools
|
|
|
|
import numpy as np
|
|
import onnx
|
|
import onnxruntime as ort
|
|
|
|
|
|
ORT_TYPES_TO_NP_TYPES = {
|
|
"tensor(float16)": np.float16,
|
|
"tensor(float)": np.float32,
|
|
"tensor(uint8)": np.uint8,
|
|
}
|
|
|
|
|
|
def _promote_raw_half_blob(attribute):
|
|
float32_values = np.frombuffer(attribute.raw_data, dtype=np.float16)
|
|
attribute.data_type = 1
|
|
attribute.raw_data = float32_values.astype(np.float32).tobytes()
|
|
|
|
|
|
def _rewrite_tensor_io_types(model):
|
|
for value_info in itertools.chain(model.graph.input, model.graph.output):
|
|
if value_info.type.tensor_type.elem_type == 10:
|
|
value_info.type.tensor_type.elem_type = 1
|
|
|
|
|
|
def _rewrite_cast_nodes(model):
|
|
for node in model.graph.node:
|
|
if node.op_type == "Cast" and node.attribute[0].i == 10:
|
|
node.attribute[0].i = 1
|
|
for attribute in node.attribute:
|
|
if hasattr(attribute, "t") and attribute.t.data_type == 10:
|
|
_promote_raw_half_blob(attribute.t)
|
|
|
|
|
|
def attributeproto_fp16_to_fp32(attr):
|
|
_promote_raw_half_blob(attr)
|
|
|
|
|
|
def convert_fp16_to_fp32(model):
|
|
for initializer in model.graph.initializer:
|
|
if initializer.data_type == 10:
|
|
_promote_raw_half_blob(initializer)
|
|
_rewrite_tensor_io_types(model)
|
|
_rewrite_cast_nodes(model)
|
|
return model.SerializeToString()
|
|
|
|
|
|
def _cpu_session_options():
|
|
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
|
|
return options
|
|
|
|
|
|
def make_onnx_cpu_runner(model_path):
|
|
model_blob = convert_fp16_to_fp32(onnx.load(model_path))
|
|
return ort.InferenceSession(model_blob, _cpu_session_options(), providers=["CPUExecutionProvider"])
|