Skip to content

Commit 28f3d18

Browse files
jamjamjonjamjamjon
authored
Bump ort from 2.0.0-rc.9 to 2.0.0-rc.10 (#107)
1 parent a3a4bf4 commit 28f3d18

9 files changed

Lines changed: 291 additions & 169 deletions

File tree

‎.github/workflows/rust-ci.yml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ jobs:
2222
- name: Install dependencies
2323
run: |
2424
DEBIAN_FRONTEND=noninteractive apt-get update --fix-missing
25-
DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential ca-certificates clang curl pkg-config protobuf-compiler
25+
DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential libssl-dev ca-certificates clang curl pkg-config protobuf-compiler
2626
2727
- name: Setup Rust
2828
uses: dtolnay/rust-toolchain@stable
@@ -47,7 +47,7 @@ jobs:
4747
- name: Install dependencies
4848
run: |
4949
DEBIAN_FRONTEND=noninteractive apt-get update --fix-missing
50-
DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential ca-certificates clang curl pkg-config protobuf-compiler
50+
DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential libssl-dev ca-certificates clang curl pkg-config protobuf-compiler
5151
5252
- name: Setup Rust
5353
uses: dtolnay/rust-toolchain@stable
@@ -67,7 +67,7 @@ jobs:
6767
- name: Install dependencies
6868
run: |
6969
DEBIAN_FRONTEND=noninteractive apt-get update --fix-missing
70-
DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential ca-certificates clang curl pkg-config protobuf-compiler
70+
DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential libssl-dev ca-certificates clang curl pkg-config protobuf-compiler
7171
7272
- name: Setup Rust
7373
uses: dtolnay/rust-toolchain@nightly
@@ -93,7 +93,7 @@ jobs:
9393
- name: Install dependencies
9494
run: |
9595
DEBIAN_FRONTEND=noninteractive apt-get update --fix-missing
96-
DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential ca-certificates clang curl pkg-config protobuf-compiler
96+
DEBIAN_FRONTEND=noninteractive apt-get install -y build-essential libssl-dev ca-certificates clang curl pkg-config protobuf-compiler
9797
9898
- name: Setup Rust
9999
uses: dtolnay/rust-toolchain@stable

‎Cargo.toml‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,10 @@ fast_image_resize = { version = "5.1.2", features = ["image"] }
3838
ndarray-npy = "0.9.1"
3939
half = { version = "2.3.1" }
4040
prost = "0.13.5"
41-
ort = { version = "2.0.0-rc.9", default-features = false, optional = true, features = [
42-
"ndarray",
41+
ort = { version = "=2.0.0-rc.10", default-features = false, optional = true, features = [
4342
"copy-dylibs",
4443
"half",
44+
"std",
4545
] }
4646
tokenizers = { version = "0.21.1" }
4747
paste = "1.0.15"
@@ -54,10 +54,10 @@ argh = "0.1.13"
5454
tracing-subscriber = { version = "0.3.18", features = ["env-filter", "chrono"] }
5555

5656
[features]
57-
default = ["ort-download-binaries"]
58-
ort-download-binaries = ["ort", "ort/download-binaries"]
59-
ort-load-dynamic = ["ort", "ort/load-dynamic"]
60-
cuda = ["ort/cuda"]
61-
trt = ["ort/tensorrt"]
62-
mps = ["ort/coreml"]
63-
video = ["dep:video-rs"]
57+
default = [ "ort-download-binaries" ]
58+
video = [ "dep:video-rs" ]
59+
ort-download-binaries = [ "ort", "ort/download-binaries" ]
60+
ort-load-dynamic = [ "ort", "ort/load-dynamic" ]
61+
cuda = [ "ort/cuda" ]
62+
trt = [ "ort/tensorrt" ]
63+
coreml = [ "ort/coreml" ]

‎build.rs‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,5 @@ use std::io::Result;
33
fn main() -> Result<()> {
44
prost_build::compile_protos(&["src/utils/onnx.proto3"], &["src"])?;
55

6-
#[cfg(any(target_os = "macos", target_os = "ios", target_os = "tvos"))]
7-
println!("cargo:rustc-link-arg=-fapple-link-rtlib");
8-
96
Ok(())
107
}

‎src/inference/engine.rs‎

Lines changed: 65 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,12 @@ use crate::{
2020
impl From<TensorElementType> for DType {
2121
fn from(dtype: TensorElementType) -> Self {
2222
match dtype {
23+
TensorElementType::Int4 => Self::Int4,
2324
TensorElementType::Int8 => Self::Int8,
2425
TensorElementType::Int16 => Self::Int16,
2526
TensorElementType::Int32 => Self::Int32,
2627
TensorElementType::Int64 => Self::Int64,
28+
TensorElementType::Uint4 => Self::Uint4,
2729
TensorElementType::Uint8 => Self::Uint8,
2830
TensorElementType::Uint16 => Self::Uint16,
2931
TensorElementType::Uint32 => Self::Uint32,
@@ -32,14 +34,19 @@ impl From<TensorElementType> for DType {
3234
TensorElementType::Float32 => Self::Fp32,
3335
TensorElementType::Float64 => Self::Fp64,
3436
TensorElementType::Bfloat16 => Self::Bf16,
35-
TensorElementType::String => Self::String,
36-
TensorElementType::Bool => Self::Bool,
37+
TensorElementType::Float8E4M3FN => Self::Fp8e4m3fn,
38+
TensorElementType::Float8E4M3FNUZ => Self::Fp8e4m3fnuz,
39+
TensorElementType::Float8E5M2 => Self::Fp8e5m2,
40+
TensorElementType::Float8E5M2FNUZ => Self::Fp8e5m2fnuz,
41+
TensorElementType::Complex64 => Self::Complex64,
42+
TensorElementType::Complex128 => Self::Complex128,
43+
_ => todo!(),
3744
}
3845
}
3946
}
4047

4148
/// A struct for tensor attrs composed of the names, the dtypes, and the dimensions.
42-
#[derive(Builder, Debug, Clone)]
49+
#[derive(Builder, Debug, Clone, Default)]
4350
pub struct OrtTensorAttr {
4451
pub names: Vec<String>,
4552
pub dtypes: Vec<TensorElementType>,
@@ -133,7 +140,9 @@ impl Engine {
133140
let param = tensor_proto.dims.iter().product::<i64>() as usize;
134141
params += param;
135142
let param = Ops::make_divisible(param, byte_alignment);
136-
let n = Self::nbytes_from_onnx_dtype_id(tensor_proto.data_type as usize);
143+
let n = Self::get_ort_dtype_from_proto_dtype_id(tensor_proto.data_type)
144+
.map(|x| x.byte_size(1))
145+
.unwrap_or_default();
137146
let wbmem = param * n;
138147
wbmems += wbmem;
139148
}
@@ -145,7 +154,10 @@ impl Engine {
145154
let param = tensor.dims.iter().product::<i64>() as usize;
146155
params += param;
147156
let param = Ops::make_divisible(param, byte_alignment);
148-
let n = Self::nbytes_from_onnx_dtype_id(tensor.data_type as usize);
157+
let n = Self::get_ort_dtype_from_proto_dtype_id(tensor.data_type)
158+
.map(|x| x.byte_size(1))
159+
.unwrap_or_default();
160+
149161
let wbmem = param * n;
150162
wbmems += wbmem;
151163
}
@@ -211,7 +223,7 @@ impl Engine {
211223

212224
// update
213225
pb.set_message(format!(
214-
"{}({}) on {:?}",
226+
"{}({}) on {}",
215227
self.spec,
216228
match self.params {
217229
Some(bytes) if bytes != 0 => {
@@ -231,7 +243,7 @@ impl Engine {
231243

232244
pub fn run(&mut self, xs: Xs) -> Result<Xs> {
233245
let mut ys = xs.derive();
234-
if let Some(onnx) = &self.onnx {
246+
if let Some(onnx) = &mut self.onnx {
235247
// alignment
236248
let xs_ = elapsed!(&format!("[{}] ort_preprocessing", self.spec), self.ts, {
237249
let mut xs_ = Vec::new();
@@ -267,38 +279,22 @@ impl Engine {
267279

268280
fn preprocess(x: &X, dtype: &TensorElementType) -> Result<DynValue> {
269281
let x = match dtype {
270-
TensorElementType::Float32 => Value::from_array(x.view())?.into_dyn(),
271-
TensorElementType::Float16 => {
272-
Value::from_array(x.mapv(f16::from_f32).view())?.into_dyn()
273-
}
274-
TensorElementType::Float64 => Value::from_array(x.view())?.into_dyn(),
275-
TensorElementType::Bfloat16 => {
276-
Value::from_array(x.mapv(bf16::from_f32).view())?.into_dyn()
277-
}
278-
TensorElementType::Int8 => Value::from_array(x.mapv(|x_| x_ as i8).view())?.into_dyn(),
279-
TensorElementType::Int16 => {
280-
Value::from_array(x.mapv(|x_| x_ as i16).view())?.into_dyn()
281-
}
282-
TensorElementType::Int32 => {
283-
Value::from_array(x.mapv(|x_| x_ as i32).view())?.into_dyn()
284-
}
285-
TensorElementType::Int64 => {
286-
Value::from_array(x.mapv(|x_| x_ as i64).view())?.into_dyn()
282+
TensorElementType::Float32 | TensorElementType::Float64 => {
283+
Value::from_array(x.0.clone())?.into_dyn()
287284
}
288-
TensorElementType::Uint8 => Value::from_array(x.mapv(|x_| x_ as u8).view())?.into_dyn(),
289-
TensorElementType::Uint16 => {
290-
Value::from_array(x.mapv(|x_| x_ as u16).view())?.into_dyn()
291-
}
292-
TensorElementType::Uint32 => {
293-
Value::from_array(x.mapv(|x_| x_ as u32).view())?.into_dyn()
294-
}
295-
TensorElementType::Uint64 => {
296-
Value::from_array(x.mapv(|x_| x_ as u64).view())?.into_dyn()
297-
}
298-
TensorElementType::Bool => Value::from_array(x.mapv(|x_| x_ != 0.).view())?.into_dyn(),
285+
TensorElementType::Float16 => Value::from_array(x.mapv(f16::from_f32))?.into_dyn(),
286+
TensorElementType::Bfloat16 => Value::from_array(x.mapv(bf16::from_f32))?.into_dyn(),
287+
TensorElementType::Int8 => Value::from_array(x.mapv(|x_| x_ as i8))?.into_dyn(),
288+
TensorElementType::Int16 => Value::from_array(x.mapv(|x_| x_ as i16))?.into_dyn(),
289+
TensorElementType::Int32 => Value::from_array(x.mapv(|x_| x_ as i32))?.into_dyn(),
290+
TensorElementType::Int64 => Value::from_array(x.mapv(|x_| x_ as i64))?.into_dyn(),
291+
TensorElementType::Uint8 => Value::from_array(x.mapv(|x_| x_ as u8))?.into_dyn(),
292+
TensorElementType::Uint16 => Value::from_array(x.mapv(|x_| x_ as u16))?.into_dyn(),
293+
TensorElementType::Uint32 => Value::from_array(x.mapv(|x_| x_ as u32))?.into_dyn(),
294+
TensorElementType::Uint64 => Value::from_array(x.mapv(|x_| x_ as u64))?.into_dyn(),
295+
TensorElementType::Bool => Value::from_array(x.mapv(|x_| x_ != 0.))?.into_dyn(),
299296
_ => unimplemented!(),
300297
};
301-
302298
Ok(x)
303299
}
304300

@@ -307,7 +303,7 @@ impl Engine {
307303
where
308304
T: Clone + 'static + ort::tensor::PrimitiveTensorElementType,
309305
{
310-
match x.try_extract_tensor::<T>() {
306+
match x.try_extract_array::<T>() {
311307
Err(err) => {
312308
debug!("Failed to extract from ort outputs: {:?}. A default value has been generated.", err);
313309
Array::zeros(0).into_dyn()
@@ -344,7 +340,7 @@ impl Engine {
344340
\nConsider enabling them by passing, e.g., `--features #FEATURE`";
345341

346342
match self.device {
347-
Device::TensorRT(id) => {
343+
Device::TensorRt(id) => {
348344
#[cfg(not(feature = "trt"))]
349345
{
350346
anyhow::bail!(feature_help
@@ -431,16 +427,28 @@ impl Engine {
431427
}
432428
}
433429
}
434-
Device::CoreML(id) => {
435-
#[cfg(not(feature = "mps"))]
430+
Device::CoreMl(id) => {
431+
#[cfg(not(feature = "coreml"))]
436432
{
437433
anyhow::bail!(feature_help
438434
.replace("#EP", "CoreML")
439-
.replace("#FEATURE", "mps"));
435+
.replace("#FEATURE", "coreml"));
440436
}
441-
#[cfg(feature = "mps")]
437+
#[cfg(feature = "coreml")]
442438
{
443-
let ep = ort::execution_providers::CoreMLExecutionProvider::default();
439+
let ep = ort::execution_providers::CoreMLExecutionProvider::default()
440+
.with_model_cache_dir(
441+
crate::Dir::Cache
442+
.crate_dir_default_with_subs(&["coreml-cache"])?
443+
.display(),
444+
)
445+
.with_compute_units(ort::execution_providers::coreml::CoreMLComputeUnits::All)
446+
.with_static_input_shapes(false)
447+
.with_subgraphs(true)
448+
.with_model_format(ort::execution_providers::coreml::CoreMLModelFormat::MLProgram)
449+
.with_specialization_strategy(
450+
ort::execution_providers::coreml::CoreMLSpecializationStrategy::FastPrediction,
451+
);
444452
match ep.is_available() {
445453
Ok(true) => {
446454
ep.register(&mut builder).map_err(|err| {
@@ -452,13 +460,14 @@ impl Engine {
452460
}
453461
}
454462
_ => {
455-
let ep = ort::execution_providers::CPUExecutionProvider::default();
463+
let ep = ort::execution_providers::CPUExecutionProvider::default()
464+
.with_arena_allocator(true);
456465
match ep.is_available() {
457466
Ok(true) => {
458467
ep.register(&mut builder)
459468
.map_err(|err| anyhow::anyhow!("Failed to register Cpu: {}", err))?;
460469
}
461-
_ => anyhow::bail!(compile_help.replace("#EP", "Cpu")),
470+
_ => unreachable!("CPU EP is not available. This case should ideally not be reached under normal circumstances."),
462471
}
463472
}
464473
}
@@ -532,38 +541,8 @@ impl Engine {
532541
Ok(ys)
533542
}
534543

535-
#[allow(dead_code)]
536-
fn nbytes_from_onnx_dtype_id(x: usize) -> usize {
537-
match x {
538-
7 | 11 | 13 => 8, // i64, f64, u64
539-
1 | 6 | 12 => 4, // f32, i32, u32
540-
10 | 16 | 5 | 4 => 2, // f16, bf16, i16, u16
541-
2 | 3 | 9 => 1, // u8, i8, bool
542-
8 => 4, // string(1~4)
543-
_ => 1, // TODO: others
544-
}
545-
}
546-
547-
#[allow(dead_code)]
548-
fn nbytes_from_onnx_dtype(x: &TensorElementType) -> usize {
549-
match x {
550-
TensorElementType::Float64 | TensorElementType::Uint64 | TensorElementType::Int64 => 8, // i64, f64, u64
551-
TensorElementType::Float32
552-
| TensorElementType::Uint32
553-
| TensorElementType::Int32
554-
| TensorElementType::String => 4, // f32, i32, u32, string(1~4)
555-
TensorElementType::Float16
556-
| TensorElementType::Bfloat16
557-
| TensorElementType::Int16
558-
| TensorElementType::Uint16 => 2, // f16, bf16, i16, u16
559-
TensorElementType::Uint8 | TensorElementType::Int8 | TensorElementType::Bool => 1, // u8, i8, bool
560-
}
561-
}
562-
563-
#[allow(dead_code)]
564-
fn ort_dtype_from_onnx_dtype_id(value: i32) -> Option<TensorElementType> {
544+
fn get_ort_dtype_from_proto_dtype_id(value: i32) -> Option<TensorElementType> {
565545
match value {
566-
0 => None,
567546
1 => Some(TensorElementType::Float32),
568547
2 => Some(TensorElementType::Uint8),
569548
3 => Some(TensorElementType::Int8),
@@ -577,10 +556,16 @@ impl Engine {
577556
11 => Some(TensorElementType::Float64),
578557
12 => Some(TensorElementType::Uint32),
579558
13 => Some(TensorElementType::Uint64),
580-
14 => None, // COMPLEX64
581-
15 => None, // COMPLEX128
559+
14 => Some(TensorElementType::Complex64),
560+
15 => Some(TensorElementType::Complex128),
582561
16 => Some(TensorElementType::Bfloat16),
583-
_ => None,
562+
17 => Some(TensorElementType::Float8E4M3FN),
563+
18 => Some(TensorElementType::Float8E4M3FNUZ),
564+
19 => Some(TensorElementType::Float8E5M2),
565+
20 => Some(TensorElementType::Float8E5M2FNUZ),
566+
21 => Some(TensorElementType::Uint4),
567+
22 => Some(TensorElementType::Int4),
568+
_ => None, // 23: Float4e2m1, 0: Undefined
584569
}
585570
}
586571

@@ -609,7 +594,7 @@ impl Engine {
609594
_ => continue,
610595
};
611596
let tensor_type = tensor.elem_type;
612-
let tensor_type = match Self::ort_dtype_from_onnx_dtype_id(tensor_type) {
597+
let tensor_type = match Self::get_ort_dtype_from_proto_dtype_id(tensor_type) {
613598
Some(dtype) => dtype,
614599
None => continue,
615600
};
@@ -642,24 +627,6 @@ impl Engine {
642627
})
643628
}
644629

645-
// pub fn to_ort(&self) -> TensorElementType {
646-
// match self {
647-
// Self::Int8 => TensorElementType::Int8,
648-
// Self::Int16 => TensorElementType::Int16,
649-
// Self::Int32 => TensorElementType::Int32,
650-
// Self::Int64 => TensorElementType::Int64,
651-
// Self::Uint8 => TensorElementType::Uint8,
652-
// Self::Uint16 => TensorElementType::Uint16,
653-
// Self::Uint32 => TensorElementType::Uint32,
654-
// Self::Uint64 => TensorElementType::Uint64,
655-
// Self::Fp16 => TensorElementType::Float16,
656-
// Self::Fp32 => TensorElementType::Float32,
657-
// Self::Fp64 => TensorElementType::Float64,
658-
// Self::Bf16 => TensorElementType::Bfloat16,
659-
// _ => todo!(),
660-
// }
661-
// }
662-
663630
pub fn load_onnx<P: AsRef<std::path::Path>>(p: P) -> Result<onnx::ModelProto> {
664631
let f = std::fs::read(p.as_ref())?;
665632
onnx::ModelProto::decode(f.as_slice()).map_err(|err| {

‎src/inference/x.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
use anyhow::Result;
22
use image::DynamicImage;
33
use ndarray::{Array, Dim, IntoDimension, Ix2, IxDyn, IxDynImpl};
4-
// use std::ops::Mul;
54

65
use crate::{Ops, ResizeMode};
76

0 commit comments

Comments
 (0)