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 软件包对增强后的模型执行推理。

用于视觉、文本和 NLP 模型的预处理和后处理自定义算子 此图片使用 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 项目中,对以下内容进行更改:

  1. build.gradle (Project)

     repositories {
         mavenCentral()
     }
    
  2. 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 图添加预处理和后处理

如果预处理算子是一个 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 联系我们以提出任何其他问题或意见。

许可证

MIT 许可证


目录