Package: onnx-extended
onnx-extended: extensions for onnx and onnxruntimeonnx-extendedextends the list of supported operators in onnx
reference implementation andonnxruntime,
or implements faster versions in C++.
Documentationonnx-extended.
Source are available ongithub/onnx-extended.Use a C++ implementation of existing operatorsimporttimeitimportnumpyasnpfromonnximportTensorProtofromonnx.helperimport(make_graph,make_model,make_node,make_opsetid,make_tensor_value_info,)fromonnx.referenceimportReferenceEvaluatorfromonnxruntimeimportInferenceSessionfromonnx_extended.ext_test_caseimportmeasure_timefromonnx_extended.referenceimportCReferenceEvaluatorX=make_tensor_value_info("X",TensorProto.FLOAT,[None,None,None,None])Y=make_tensor_value_info("Y",TensorProto.FLOAT,[None,None,None,None])B=make_tensor_value_info("B",TensorProto.FLOAT,[None,None,None,None])W=make_tensor_value_info("W",TensorProto.FLOAT,[None,None,None,None])node=make_node("Conv",["X","W","B"],["Y"],pads=[1,1,1,1],dilations=[1,1],strides=[2,2],)graph=make_graph([node],"g",[X,W,B],[Y])onnx_model=make_model(graph,opset_imports=[make_opsetid("",16)])sH,sW=64,64X=np.arange(sW*sH).reshape((1,1,sH,sW)).astype(np.float32)W=np.ones((1,1,3,3),dtype=np.float32)B=np.array([[[[0]]]],dtype=np.float32)sess1=ReferenceEvaluator(onnx_model)sess2=CReferenceEvaluator(onnx_model)# 100 times fasterexpected=sess1.run(None,{"X":X,"W":W,"B":B})[0]got=sess2.run(None,{"X":X,"W":W,"B":B})[0]diff=np.abs(expected-got).max()print(f"difference:{diff}")f1=lambda:sess1.run(None,{"X":X,"W":W,"B":B})[0]f2=lambda:sess2.run(None,{"X":X,"W":W,"B":B})[0]print("onnx:",timeit.timeit(f1,globals=globals(),number=5))print("onnx-extended:",timeit.timeit(f2,globals=globals(),number=5))difference: 0.0
onnx: 0.024006774998269975
onnx-extended: 0.0002316169993719086Build with CUDA, openmp, eigen, onnxruntimeThe package also contains some dummy examples on how to
build with C++ functions (pybind11,cython),
withopenmp,eigenwith or without CUDA. It also shows how to create a custom operator
foronnxruntimein C++.The version released onpypi/onnx-extendedonly works on CPU. It needs to be manually built to enable
the code using CUDA. The build will automatically link with CUDA if it is found.
If not, some extensions might not be available.python setup.py build_ext --inplace
# pip install -e .It is possible to use a specific version of CUDA:python setup.py build_ext --inplace --cuda-version=11.8
# or (not working yet)
# pip install -e . --config-settings="--cuda-version=11.8"
# pip install -e . --global-option="--cuda-version=11.8"
export USE_CUDA=11.8
pip install -e .NVTXcan be enabled with the following command:python setup.py build_ext --inplace --use_nvtx 1
# or (not working yet)
# pip install -e . --config-settings="--use_nvtx=1"
pip install -e . --global-option "--use_nvtx=1"Experimental cython binding for onnxruntimeThe python onnxruntime package relies on pybind11 to expose
its functionalities.onnx-extendedtries to build a cython wrapper
around the C/C++ API of onnxruntime. cython relies on python C API
and is faster than pybind11. This different may be significant when
onnxruntime is used on small graphs and tensors.Custom kernels for onnxruntimeonnxruntime provides an API to add custom implementation
for existing or new onnx operators. An example for CPU.from onnxruntime import InferenceSession, SessionOptions
from onnx_extended.ortops.optim.cpu import get_ort_ext_libs

r = get_ort_ext_libs()
opts = SessionOptions()
if r is not None:
    opts.register_custom_ops_library(r[0])

sess_cus = InferenceSession(
    onx_modified.SerializeToString(), opts, providers=["CPUExecutionProvider"]
)