ONNX Runtime 扩展
ONNX Runtime 扩展是一个通过 ONNX Runtime 自定义算子接口扩展 ONNX 模型能力及使用 ONNX Runtime 进行推理的库。它包含一组自定义算子,用于支持音频、视觉、文本和语言模型的常见模型预处理和后处理。与 ONNX Runtime 一样,Extensions 还支持多种语言和平台(Windows/Linux/macOS 上的 Python、Android 和 iOS 移动平台以及用于 Web 的 WebAssembly)。
基本工作流是将自定义算子添加到 ONNX 模型中,然后使用 ONNX Runtime 和 ONNXRuntime-Extensions 软件包对增强后的模型执行推理。
此图片使用 Combine.AI 创建,该工具由 Bing Chat、Bing Image Creator 和 EdgeGPT 提供支持。
快速入门
Python 安装
pip install onnxruntime-extensions
每日构建
在 Windows 上
pip install --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/ onnxruntime-extensions
onnxruntime-extensions 软件包依赖于 onnx 和 onnxruntime。
在 Linux/macOS 上
在执行以下命令之前,请确保已安装如 gcc(g++ 8.0 或更高版本)或 clang 等编译器工具包
python -m pip install git+https://github.com/microsoft/onnxruntime-extensions.git
NuGet 安装(使用 .NET CLI)
dotnet add package Microsoft.ML.OnnxRuntime.Extensions --version 0.8.1-alpha
iOS 安装
在你的 CocoaPods Podfile 中,添加 onnxruntime-extensions-c pod。
use_frameworks!
# onnxruntime C/C++ full package
pod 'onnxruntime-c'
# onnxruntime-extensions C/C++ package
pod 'onnxruntime-extensions-c'
运行 pod install。
Android 安装
在你的 Android Studio 项目中,对以下内容进行更改:
-
build.gradle (Project)
repositories { mavenCentral() } -
build.gradle (Module)
dependencies { // onnxruntime full package implementation 'com.microsoft.onnxruntime:onnxruntime-android:latest.release' // onnxruntime-extensions package implementation 'com.microsoft.onnxruntime:onnxruntime-extensions-android:latest.release' }
向模型添加预处理和后处理
有多种方法可以向 ONNX 图添加预处理和后处理
- 如果流水线 API 支持该模型及其预处理,请使用预处理流水线 API
- 从 PyTorch 模型导出为 ONNX
- 创建一个包含自定义算子节点的模型图的 ONNX 模型
- 如果预处理已经在 ONNX 图中,请使用 ONNX API 将预处理与 ONNX 模型组合
如果预处理算子是一个 HuggingFace 分词器,你还可以通过从 HuggingFace transformer 数据处理类进行转换,轻松获取 ONNX 处理图,如下例所示
import onnxruntime as _ort
from transformers import AutoTokenizer, GPT2Tokenizer
from onnxruntime_extensions import OrtPyFunction, gen_processing_models
# SentencePieceTokenizer
spm_hf_tokenizer = AutoTokenizer.from_pretrained("t5-base", model_max_length=512)
spm_onnx_model = OrtPyFunction(gen_processing_models(spm_hf_tokenizer, pre_kwargs={})[0])
# GPT2Tokenizer
gpt2_hf_tokenizer = GPT2Tokenizer.from_pretrained("Xenova/gpt-4", use_fast=False)
gpt2_onnx_model = OrtPyFunction(gen_processing_models(gpt2_hf_tokenizer, pre_kwargs={})[0])
有关更多信息,你可以通过以下方式查看 API
help(onnxruntime_extensions.gen_processing_models)
如果找不到我需要的自定义算子怎么办?
在此处查找我们当前支持的自定义算子。如果你没有找到所需的自定义算子,可以像这样向 ONNX Runtime Extensions 添加新的自定义算子。请注意,如果你确实添加了新算子,则必须从源码构建。
使用 ONNX Runtime 和 Extensions 进行推理
Python
针对以下语言有单独的软件包,请为构建安装它们。
import onnxruntime as _ort
from onnxruntime_extensions import get_library_path as _lib_path
so = _ort.SessionOptions()
so.register_custom_ops_library(_lib_path())
# Run the ONNXRuntime Session as per ONNXRuntime docs suggestions.
sess = _ort.InferenceSession(model, so)
sess.run (...)
C++
使用 Extensions 共享库的路径注册 Extensions。
Ort::Env env = ...;
// Note: use `wchar_t` instead of `char` for paths on Windows
const char* model_uri = "/path/to/the/model.onnx";
const char* custom_op_library_filename = "/path/to/the/onnxruntime-extensions/shared/library";
Ort::SessionOptions session_options;
// Register Extensions custom ops with the session options.
Ort::ThrowOnError(Ort::GetApi().RegisterCustomOpsLibrary_V2(static_cast<OrtSessionOptions*>(session_options),
custom_op_library_filename));
// Create a session.
Ort::Session session(env, model_uri, session_options);
通过直接调用 RegisterCustomOps 函数注册 Extensions。
Ort::Env env = ...;
// Note: use `wchar_t` instead of `char` for paths on Windows
const char* model_uri = "/path/to/the/model.onnx";
Ort::SessionOptions session_options;
// Register Extensions custom ops with the session options.
// `RegisterCustomOps` is declared in onnxruntime_extensions.h.
Ort::ThrowOnError(RegisterCustomOps(static_cast<OrtSessionOptions*>(session_options), OrtGetApiBase()));
// Create a session.
Ort::Session session(env, model_uri, session_options);
Java
var env = OrtEnvironment.getEnvironment();
var sess_opt = new OrtSession.SessionOptions();
/* Register the custom ops from onnxruntime-extensions */
sess_opt.registerCustomOpLibrary(OrtxPackage.getLibraryPath());
C#
SessionOptions options = new SessionOptions();
options.RegisterOrtExtensions();
session = new InferenceSession(model, options);
教程
查看使用我们自定义算子的一些端到端教程
参与贡献
本项目欢迎贡献和建议。大多数贡献要求你同意《贡献者许可协议》(CLA),声明你有权且确实授予我们使用你的贡献的权利。有关详细信息,请访问 https://cla.microsoft.com。
当你提交拉取请求时,CLA-bot 将自动确定你是否需要提供 CLA,并对 PR 进行相应的标记(例如添加标签、评论)。只需按照机器人提供的说明操作即可。在所有使用我们 CLA 的仓库中,你只需执行此操作一次。
本项目采用了 Microsoft Open Source Code of Conduct。有关更多信息,请参阅行为准则常见问题解答,或通过 opencode@microsoft.com 联系我们以提出任何其他问题或意见。