Source code for experimental_experiment.convert.ort_helper
import os
from onnx import ModelProto
[docs]
def append_custom_libraries(
onx: ModelProto,
options: "onnxruntime.SessionOptions", # noqa: F821
):
"""
Appends libraries implementing custom kernels.
The functions checks the opsets then add the necessary
custom libraries to the options.
:param onx: model proto
:param options: onnxruntime.SessionOptions
"""
domains = set(d.domain for d in onx.opset_import)
if "onnx_extended.ortops.optim.cuda" in domains:
from onnx_extended.ortops.optim.cuda import get_ort_ext_libs
assert os.path.exists(
get_ort_ext_libs()[0]
), f"Unable to find library {get_ort_ext_libs()[0]!r}."
options.register_custom_ops_library(get_ort_ext_libs()[0])
elif "onnx_extended.ortops.optim.cpu" in domains:
from onnx_extended.ortops.optim.cpu import get_ort_ext_libs
assert os.path.exists(
get_ort_ext_libs()[0]
), f"Unable to find library {get_ort_ext_libs()[0]!r}."
options.register_custom_ops_library(get_ort_ext_libs()[0])