2个稳定版本
1.16.1 | 2023年10月31日 |
---|---|
1.15.1 | 2023年7月23日 |
#70 in 机器学习
1.5MB
53K SLoC
ONNX Runtime
这是对Microsoft的ONNX Runtime (版本1.8)的Rust包装器的尝试。
该项目包含两个crate
onnxruntime-sys
: 到C API的低级绑定;onnxruntime
: 高级和安全的API。
build.rs 脚本支持下载Microsoft ONNX Runtime的预构建版本,提供了以下目标
CPU
- Linux x86_64
- macOS x86_64
- macOS aarch64(没有预构建二进制文件,没有CI测试,见 #74)
- Windows i686
- Windows x86_64
GPU
- Linux x86_64
- Windows x86_64
警告:
- 这是一个实验性的工作,仍在进行中;它不是完整/工作/安全的。欢迎提供帮助!
- 基本推理功能正常,见
onnxruntime/examples/sample.rs
或onnxruntime/tests/integration_tests.rs
- ONNX Runtime有许多选项来控制推理过程,但这些选项尚未公开。
- 该软件在macOS Catalina上开发和测试。其他平台应该也能工作,但尚未测试。
设置
build.rs 脚本支持三种不同的策略来获取ONNX Runtime
- 从上游下载预构建的二进制文件;
- 指向已安装的本地版本;
- 从源代码编译(尚未实现)。
要选择要使用的策略,请将环境变量 ORT_STRATEGY
设置为
download
:如果未设置ORT_STRATEGY
,则为默认值;system
:使用本地安装的版本(使用环境变量ORT_LIB_LOCATION
指向安装路径)compile
:编译库
下载策略 download
支持下载支持 CUDA 的 ONNX 版本。要使用此功能,请在 Cargo.toml
中设置特性 cuda
。
在构建脚本允许编译运行时之前,请参阅编译说明以了解有关此过程的详细信息。
CUDA 使用说明
要使用 CUDA,您需要设置特性 cuda
,并且还需要使用方法 use_cuda
设置您的会话,如下所示
let mut session = environment
.new_session_builder()?
.use_cuda(0)?
'ORT_STRATEGY=system' 说明
当使用 ORT_STRATEGY=system
时,如果在系统路径中未安装库,则在 macOS 上执行构建的二进制文件(例如测试)可能会失败。可能会发生类似于以下错误的错误
dyld: Library not loaded: @rpath/libonnxruntime.1.7.1.dylib
Referenced from: onnxruntime-rs.git/target/debug/deps/onnxruntime_sys-22eb0e3e89a0278c
Reason: image not found
要修复此问题,可以执行以下操作之一
-
将环境变量
LD_LIBRARY_PATH
设置为指向可以找到库的路径。 -
将
.cargo/config
文件调整为包含链接器标志以提供 完整 路径[target.aarch64-apple-darwin] rustflags = ["-C", "link-args=-Wl,-rpath,/full/path/to/onnxruntime/lib"]
有关更多信息,请参阅 rust-lang/cargo #5077。
示例
使用 C API 的 C++ 示例(C_Api_Sample.cpp
)被移植到低级别包(onnxruntime-sys
)和高级别包(onnxruntime
)。
onnxruntime-sys
运行此示例(onnxruntime-sys/examples/c_api_sample.rs
)
# Download the model (SqueezeNet 1.0, ONNX version: 1.3, Opset version: 8)
❯ curl -LO "https://github.com/onnx/models/raw/master/vision/classification/squeezenet/model/squeezenet1.0-8.onnx"
❯ cargo run --example c_api_sample
[...]
Finished dev [unoptimized + debuginfo] target(s) in 1.88s
Running `target/debug/examples/c_api_sample`
Using Onnxruntime C API
2020-08-09 09:37:41.554922 [I:onnxruntime:, inference_session.cc:174 ConstructorCommon] Creating and using per session threadpools since use_per_session_threads_ is true
2020-08-09 09:37:41.556650 [I:onnxruntime:, inference_session.cc:830 Initialize] Initializing session.
2020-08-09 09:37:41.556665 [I:onnxruntime:, inference_session.cc:848 Initialize] Adding default CPU execution provider.
2020-08-09 09:37:41.556678 [I:onnxruntime:test, bfc_arena.cc:15 BFCArena] Creating BFCArena for Cpu
2020-08-09 09:37:41.556687 [V:onnxruntime:test, bfc_arena.cc:32 BFCArena] Creating 21 bins of max chunk size 256 to 268435456
2020-08-09 09:37:41.558313 [I:onnxruntime:, reshape_fusion.cc:37 ApplyImpl] Total fused reshape node count: 0
2020-08-09 09:37:41.559327 [I:onnxruntime:, reshape_fusion.cc:37 ApplyImpl] Total fused reshape node count: 0
2020-08-09 09:37:41.559476 [I:onnxruntime:, reshape_fusion.cc:37 ApplyImpl] Total fused reshape node count: 0
2020-08-09 09:37:41.559607 [V:onnxruntime:, inference_session.cc:671 TransformGraph] Node placements
2020-08-09 09:37:41.559615 [V:onnxruntime:, inference_session.cc:673 TransformGraph] All nodes have been placed on [CPUExecutionProvider].
2020-08-09 09:37:41.559639 [I:onnxruntime:, session_state.cc:25 SetGraph] SaveMLValueNameIndexMapping
2020-08-09 09:37:41.559787 [I:onnxruntime:, session_state.cc:70 SetGraph] Done saving OrtValue mappings.
2020-08-09 09:37:41.560252 [I:onnxruntime:, session_state_initializer.cc:178 SaveInitializedTensors] Saving initialized tensors.
2020-08-09 09:37:41.563467 [I:onnxruntime:, session_state_initializer.cc:223 SaveInitializedTensors] Done saving initialized tensors
2020-08-09 09:37:41.563979 [I:onnxruntime:, inference_session.cc:919 Initialize] Session successfully initialized.
Number of inputs = 1
Input 0 : name=data_0
Input 0 : type=1
Input 0 : num_dims=4
Input 0 : dim 0=1
Input 0 : dim 1=3
Input 0 : dim 2=224
Input 0 : dim 3=224
2020-08-09 09:37:41.573127 [I:onnxruntime:, sequential_executor.cc:145 Execute] Begin execution
2020-08-09 09:37:41.573183 [I:onnxruntime:test, bfc_arena.cc:259 AllocateRawInternal] Extending BFCArena for Cpu. bin_num:13 rounded_bytes:3154176
2020-08-09 09:37:41.573197 [I:onnxruntime:test, bfc_arena.cc:143 Extend] Extended allocation by 4194304 bytes.
2020-08-09 09:37:41.573203 [I:onnxruntime:test, bfc_arena.cc:147 Extend] Total allocated bytes: 9137152
2020-08-09 09:37:41.573212 [I:onnxruntime:test, bfc_arena.cc:150 Extend] Allocated memory at 0x7fb7d6cb7000 to 0x7fb7d70b7000
2020-08-09 09:37:41.573248 [I:onnxruntime:test, bfc_arena.cc:259 AllocateRawInternal] Extending BFCArena for Cpu. bin_num:8 rounded_bytes:65536
2020-08-09 09:37:41.573256 [I:onnxruntime:test, bfc_arena.cc:143 Extend] Extended allocation by 4194304 bytes.
2020-08-09 09:37:41.573262 [I:onnxruntime:test, bfc_arena.cc:147 Extend] Total allocated bytes: 13331456
2020-08-09 09:37:41.573268 [I:onnxruntime:test, bfc_arena.cc:150 Extend] Allocated memory at 0x7fb7d70b7000 to 0x7fb7d74b7000
Score for class [0] = 0.000045440644
Score for class [1] = 0.0038458651
Score for class [2] = 0.00012494653
Score for class [3] = 0.0011804523
Score for class [4] = 0.0013169361
Done!
onnxruntime
运行此示例(onnxruntime/examples/sample.rs
)
# Download the model (SqueezeNet 1.0, ONNX version: 1.3, Opset version: 8)
❯ curl -LO "https://github.com/onnx/models/raw/master/vision/classification/squeezenet/model/squeezenet1.0-8.onnx"
❯ cargo run --example sample
[...]
Finished dev [unoptimized + debuginfo] target(s) in 13.62s
Running `target/debug/examples/sample`
Uninitialized environment found, initializing it with name "test".
2020-08-09 09:34:37.395577 [I:onnxruntime:, inference_session.cc:174 ConstructorCommon] Creating and using per session threadpools since use_per_session_threads_ is true
2020-08-09 09:34:37.399253 [I:onnxruntime:, inference_session.cc:830 Initialize] Initializing session.
2020-08-09 09:34:37.399284 [I:onnxruntime:, inference_session.cc:848 Initialize] Adding default CPU execution provider.
2020-08-09 09:34:37.399313 [I:onnxruntime:test, bfc_arena.cc:15 BFCArena] Creating BFCArena for Cpu
2020-08-09 09:34:37.399335 [V:onnxruntime:test, bfc_arena.cc:32 BFCArena] Creating 21 bins of max chunk size 256 to 268435456
2020-08-09 09:34:37.410516 [I:onnxruntime:, reshape_fusion.cc:37 ApplyImpl] Total fused reshape node count: 0
2020-08-09 09:34:37.417478 [I:onnxruntime:, reshape_fusion.cc:37 ApplyImpl] Total fused reshape node count: 0
2020-08-09 09:34:37.420131 [I:onnxruntime:, reshape_fusion.cc:37 ApplyImpl] Total fused reshape node count: 0
2020-08-09 09:34:37.422623 [V:onnxruntime:, inference_session.cc:671 TransformGraph] Node placements
2020-08-09 09:34:37.428863 [V:onnxruntime:, inference_session.cc:673 TransformGraph] All nodes have been placed on [CPUExecutionProvider].
2020-08-09 09:34:37.428954 [I:onnxruntime:, session_state.cc:25 SetGraph] SaveMLValueNameIndexMapping
2020-08-09 09:34:37.429079 [I:onnxruntime:, session_state.cc:70 SetGraph] Done saving OrtValue mappings.
2020-08-09 09:34:37.429925 [I:onnxruntime:, session_state_initializer.cc:178 SaveInitializedTensors] Saving initialized tensors.
2020-08-09 09:34:37.436300 [I:onnxruntime:, session_state_initializer.cc:223 SaveInitializedTensors] Done saving initialized tensors
2020-08-09 09:34:37.437255 [I:onnxruntime:, inference_session.cc:919 Initialize] Session successfully initialized.
Dropping the session options.
2020-08-09 09:34:37.448956 [I:onnxruntime:, sequential_executor.cc:145 Execute] Begin execution
2020-08-09 09:34:37.449041 [I:onnxruntime:test, bfc_arena.cc:259 AllocateRawInternal] Extending BFCArena for Cpu. bin_num:13 rounded_bytes:3154176
2020-08-09 09:34:37.449072 [I:onnxruntime:test, bfc_arena.cc:143 Extend] Extended allocation by 4194304 bytes.
2020-08-09 09:34:37.449087 [I:onnxruntime:test, bfc_arena.cc:147 Extend] Total allocated bytes: 9137152
2020-08-09 09:34:37.449104 [I:onnxruntime:test, bfc_arena.cc:150 Extend] Allocated memory at 0x7fb3b9585000 to 0x7fb3b9985000
2020-08-09 09:34:37.449176 [I:onnxruntime:test, bfc_arena.cc:259 AllocateRawInternal] Extending BFCArena for Cpu. bin_num:8 rounded_bytes:65536
2020-08-09 09:34:37.449196 [I:onnxruntime:test, bfc_arena.cc:143 Extend] Extended allocation by 4194304 bytes.
2020-08-09 09:34:37.449209 [I:onnxruntime:test, bfc_arena.cc:147 Extend] Total allocated bytes: 13331456
2020-08-09 09:34:37.449222 [I:onnxruntime:test, bfc_arena.cc:150 Extend] Allocated memory at 0x7fb3b9985000 to 0x7fb3b9d85000
Dropping Tensor.
Score for class [0] = 0.000045440578
Score for class [1] = 0.0038458686
Score for class [2] = 0.0001249467
Score for class [3] = 0.0011804511
Score for class [4] = 0.00131694
Dropping TensorFromOrt.
Dropping the session.
Dropping the memory information.
Dropping the environment.
还可以参阅集成测试(onnxruntime/tests/integration_tests.rs
),该测试执行简单的模型下载和推理,验证结果。
绑定生成
绑定(onnxruntime-sys
的基础)已提交到 Git 存储库。这意味着 bindgen
不再是每个构建的依赖项(已使其可选),因此构建时间更好。
要生成新的绑定(例如,如果您的平台没有这些绑定,或者发生了版本升级),请使用带有 generate-bindings
特性的 crate 构建。
注意:请确保具有 rustfmt
rustup 组件,以便格式化绑定
rustup component add rustfmt
然后在每个平台上使用适当的特性标志构建
❯ cd onnxruntime-sys-ng
❯ cargo build --features generate-bindings
使用 Docker 在 Linux 上生成绑定
准备容器
❯ docker run -it --rm --name rustbuilder -v "$PWD":/usr/src/myapp -w /usr/src/myapp rust:1.50.0 /bin/bash
❯ apt-get update
❯ apt-get install clang
❯ rustup component add rustfmt
生成绑定
❯ docker exec -it --user "$(id -u)":"$(id -g)" rustbuilder /bin/bash
❯ cd onnxruntime-sys-ng
❯ cargo build --features 'generate-bindings, cuda'
使用 Vagrant 在 Windows 上生成绑定
您可以使用 nbigaouette/windows_vagrant_rust 来配置一个 Windows 虚拟机,该虚拟机可以构建项目并生成绑定。
Windows 可以构建 x86 和 x86_64 绑定
❯ rustup target add i686-pc-windows-msvc x86_64-pc-windows-msvc
❯ cd onnxruntime-sys-ng
❯ cargo build --features 'generate-bindings, cuda' --target i686-pc-windows-msvc
❯ cargo build --features 'generate-bindings, cuda' --target x86_64-pc-windows-msvc
进行
应遵守 Rust 行为准则。对于升级或管理问题,请联系 Nicolas([email protected]),而不是 Rust 管理团队。
许可证
本项目根据以下许可证之一授权
- Apache License, Version 2.0,(LICENSE-APACHE 或 https://apache.ac.cn/licenses/LICENSE-2.0)
- 麻省理工学院许可协议(LICENSE-MIT 或 http://opensource.org/licenses/MIT)
任选。
依赖项
~4.5–7MB
~148K SLoC