ONNXRuntime扩展功能深度解析：增强模型转换与推理能力-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00322/article/details/148375575

ONNXRuntime扩展功能深度解析：增强模型转换与推理能力

onnxruntime microsoft/onnxruntime: 是一个用于运行各种机器学习模型的开源库。适合对机器学习和深度学习有兴趣的人，特别是在开发和部署机器学习模型时需要处理各种不同框架和算子的人。特点是支持多种机器学习框架和算子，包括 TensorFlow、PyTorch、Caffe 等，具有高性能和广泛的兼容性。项目地址: https://gitcode.com/gh_mirrors/on/onnxruntime

一、ONNXRuntime扩展功能概述

ONNXRuntime扩展包是一个功能强大的工具集，旨在扩展ONNX模型的转换和推理能力。它为开发者提供了一系列增强功能，特别是在文本处理领域，通过引入自定义运算符大大提升了ONNX模型的表达能力。

二、核心功能：自定义运算符支持

扩展包中最具价值的部分是其丰富的自定义运算符库，主要包括两大类：

字符串运算符：处理文本数据的各种操作
分词器(Tokenizers)：支持多种流行的分词算法

这些自定义运算符极大地简化了自然语言处理(NLP)任务的实现流程，使得开发者能够更高效地构建端到端的文本处理管道。

三、构建集成扩展的ONNXRuntime

3.1 基本构建方法

要将扩展功能集成到ONNXRuntime中，开发者需要使用--use_extensions参数进行构建。这一参数会自动获取并编译扩展包为静态库，然后链接到主运行时中。

对于高级使用场景，开发者还可以：

使用预下载的扩展包路径：通过--extensions_overridden_path指定
构建最小化运行时：必须配合--minimal_build custom_ops参数

3.2 运算符配置构建

更精细的控制方式是通过运算符配置文件进行选择性构建：

# 示例配置文件
ai.onnx;12;Add,Cast,Concat,Squeeze
ai.onnx.contrib;1;GPT2Tokenizer,

配置说明：

ai.onnx.contrib是扩展运算符的专用域
每行指定了域、操作集版本和所需运算符列表

3.3 生成运算符配置

配置文件的生成有多种途径：

从现有ONNX模型自动生成
手动编辑创建（需了解模型所需的具体运算符）
使用--custom_op_library参数指定扩展库路径后生成

四、异常处理与特殊构建

在特定环境下（如WebAssembly），可能需要特殊处理异常机制：

--disable_exceptions：禁用所有异常
--enable_wasm_exception_throwing_override：针对WASM环境启用异常抛出支持

五、端到端应用示例

5.1 创建集成模型

通过ONNX帮助函数可以轻松创建包含自定义运算符的端到端模型：

# 创建分词器节点示例
tokenizer_node = helper.make_node(
    'GPT2Tokenizer',
    inputs=['input_str'],
    outputs=['token_ids', 'attention_mask'],
    vocab=vocab_content,
    merges=merges_content,
    domain='ai.onnx.contrib'
)

5.2 Python环境运行

在Python中使用扩展模型需要注册自定义运算符库：

so = _ort.SessionOptions()
so.register_custom_ops_library(_lib_path())
sess = _ort.InferenceSession(model_path, so)

5.3 JavaScript环境运行

在Web环境中使用扩展模型的基本流程：

async function runModel() {
    const session = await ort.InferenceSession.create('model.onnx');
    const results = await session.run(inputs);
    // 处理结果...
}