eliott/ort
1
0
Fork 0
mirror of https://github.com/pykeio/ort synced 2026-04-25 16:34:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: custom ops & mutable tensors

Browse Source
This commit is contained in:
Carson M.
2024-02-22 19:53:51 -06:00
parent 91463cddbc
commit 6e7c051a7a
26 changed files with 1040 additions and 334 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
tests/data/custom_op_test.onnx Normal file
View File

@@ -0,0 +1,18 @@
 OnnxRuntime:Č
8
input_1
input_2output_1" CustomOpOne:
test.customop
.
output_1output" CustomOpTwo:
test.customop CustomOpTestZ
input_1


Z
input_2


b
output


11
tests/mnist.rs
View File

@@ -1,7 +1,7 @@
use std::path::Path;
use image::{imageops::FilterType, ImageBuffer, Luma, Pixel};
use ort::{inputs, ArrayExtensions, GraphOptimizationLevel, Session, Tensor};
use ort::{inputs, ArrayExtensions, GraphOptimizationLevel, Session};
use test_log::test;
#[test]
@@ -43,8 +43,13 @@ fn mnist_5() -> ort::Result<()> {
// Perform the inference
let outputs = session.run(inputs![array]?)?;
let output: Tensor<_> = outputs[0].extract_tensor()?;
let mut probabilities: Vec<(usize, f32)> = output.view().softmax(ndarray::Axis(1)).iter().copied().enumerate().collect::<Vec<_>>();
let mut probabilities: Vec<(usize, f32)> = outputs[0]
.extract_tensor()?
.softmax(ndarray::Axis(1))
.iter()
.copied()
.enumerate()
.collect::<Vec<_>>();
// Sort probabilities so highest is at beginning of vector.
probabilities.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

11
tests/squeezenet.rs
View File

@@ -7,7 +7,7 @@ use std::{
use image::{imageops::FilterType, ImageBuffer, Pixel, Rgb};
use ndarray::s;
use ort::{inputs, ArrayExtensions, FetchModelError, GraphOptimizationLevel, Session, Tensor};
use ort::{inputs, ArrayExtensions, FetchModelError, GraphOptimizationLevel, Session};
use test_log::test;
#[test]
@@ -70,8 +70,13 @@ fn squeezenet_mushroom() -> ort::Result<()> {
// Downloaded model does not have a softmax as final layer; call softmax on second axis
// and iterate on resulting probabilities, creating an index to later access labels.
let output: Tensor<_> = outputs[0].extract_tensor()?;
let mut probabilities: Vec<(usize, f32)> = output.view().softmax(ndarray::Axis(1)).iter().copied().enumerate().collect::<Vec<_>>();
let mut probabilities: Vec<(usize, f32)> = outputs[0]
.extract_tensor()?
.softmax(ndarray::Axis(1))
.iter()
.copied()
.enumerate()
.collect::<Vec<_>>();
// Sort probabilities so highest is at beginning of vector.
probabilities.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

12
tests/upsample.rs
View File

@@ -1,8 +1,8 @@
use std::path::Path;
use image::RgbImage;
use ndarray::{Array, CowArray, Ix4};
use ort::{inputs, GraphOptimizationLevel, Session, Tensor};
use ndarray::{Array, ArrayViewD, CowArray, Ix4};
use ort::{inputs, GraphOptimizationLevel, Session};
use test_log::test;
fn load_input_image<P: AsRef<Path>>(name: P) -> RgbImage {
@@ -69,10 +69,10 @@ fn upsample() -> ort::Result<()> {
let outputs = session.run(inputs![&array]?)?;
assert_eq!(outputs.len(), 1);
let output: Tensor<f32> = outputs[0].extract_tensor()?;
let output: ArrayViewD<f32> = outputs[0].extract_tensor()?;
// The image should have doubled in size
assert_eq!(output.view().shape(), [1, 448, 448, 3]);
assert_eq!(output.shape(), [1, 448, 448, 3]);
Ok(())
}
@@ -106,10 +106,10 @@ fn upsample_with_ort_model() -> ort::Result<()> {
let outputs = session.run(inputs![&array]?)?;
assert_eq!(outputs.len(), 1);
let output: Tensor<f32> = outputs[0].extract_tensor()?;
let output: ArrayViewD<f32> = outputs[0].extract_tensor()?;
// The image should have doubled in size
assert_eq!(output.view().shape(), [1, 448, 448, 3]);
assert_eq!(output.shape(), [1, 448, 448, 3]);
Ok(())
}

7
tests/vectorizer.rs
View File

@@ -26,12 +26,7 @@ fn vectorizer() -> ort::Result<()> {
// Perform the inference
let outputs = session.run(input_tensor_values)?;
assert_eq!(
*outputs[0].extract_tensor::<f32>()?.view(),
ArrayD::from_shape_vec(IxDyn(&[1, 9]), vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
.unwrap()
.view()
);
assert_eq!(outputs[0].extract_tensor::<f32>()?, ArrayD::from_shape_vec(IxDyn(&[1, 9]), vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]).unwrap());
Ok(())
}