Merge pull request #56 from ccbrown/dynamic-size-fix

fix dynamically sized outputs

Author: nbigaouette

onnxruntime/src/error.rs

@@ -62,6 +62,9 @@ pub enum OrtError {
     /// Error occurred when checking if ONNX tensor was properly initialized
     #[error("Failed to check if tensor: {0}")]
     IsTensor(OrtApiError),
+    /// Error occurred when getting tensor type and shape
+    #[error("Failed to get tensor type and shape: {0}")]
+    GetTensorTypeAndShape(OrtApiError),
     /// Error occurred when ONNX inference operation was called
     #[error("Failed to run: {0}")]
     Run(OrtApiError),

onnxruntime/src/session.rs

@@ -405,32 +405,8 @@ impl<'a> Session<'a> {
             .map(|n| n.as_ptr() as *const i8)
             .collect();
 
-        let output_shapes: Vec<Vec<usize>> = {
-            let mut tmp = Vec::new();
-            for (idx, output) in self.outputs.iter().enumerate() {
-                let v: Vec<_> = output
-                    .dimensions
-                    .iter()
-                    .enumerate()
-                    .map(|(jdx, dim)| match dim {
-                        None => input_arrays[idx].shape()[jdx],
-                        Some(d) => *d as usize,
-                    })
-                    .collect();
-                tmp.push(v);
-            }
-            tmp
-        };
-        let memory_info_ref = &self.memory_info;
-        let output_tensor_extractors: Vec<OrtOwnedTensorExtractor<ndarray::IxDyn>> = output_shapes
-            .iter()
-            .map(|output_shape| {
-                OrtOwnedTensorExtractor::new(memory_info_ref, ndarray::IxDyn(output_shape))
-            })
-            .collect();
-
         let mut output_tensor_extractors_ptrs: Vec<*mut sys::OrtValue> =
-            vec![std::ptr::null_mut(); output_tensor_extractors.len()];
+            vec![std::ptr::null_mut(); self.outputs.len()];
 
         // The C API expects pointers for the arrays (pointers to C-arrays)
         let input_ort_tensors: Vec<OrtTensor<TIn, D>> = input_arrays
@@ -458,11 +434,23 @@ impl<'a> Session<'a> {
         };
         status_to_result(status).map_err(OrtError::Run)?;
 
+        let memory_info_ref = &self.memory_info;
         let outputs: Result<Vec<OrtOwnedTensor<TOut, ndarray::Dim<ndarray::IxDynImpl>>>> =
-            output_tensor_extractors
+            output_tensor_extractors_ptrs
                 .into_iter()
-                .zip(output_tensor_extractors_ptrs.into_iter())
-                .map(|(mut output_tensor_extractor, ptr)| {
+                .map(|ptr| {
+                    let mut tensor_info_ptr: *mut sys::OrtTensorTypeAndShapeInfo =
+                        std::ptr::null_mut();
+                    let status = unsafe {
+                        g_ort().GetTensorTypeAndShape.unwrap()(ptr, &mut tensor_info_ptr as _)
+                    };
+                    status_to_result(status).map_err(OrtError::GetTensorTypeAndShape)?;
+                    let dims = unsafe { get_tensor_dimensions(tensor_info_ptr) };
+                    unsafe { g_ort().ReleaseTensorTypeAndShapeInfo.unwrap()(tensor_info_ptr) };
+                    let dims: Vec<_> = dims?.iter().map(|&n| n as usize).collect();
+
+                    let mut output_tensor_extractor =
+                        OrtOwnedTensorExtractor::new(memory_info_ref, ndarray::IxDyn(&dims));
                     output_tensor_extractor.tensor_ptr = ptr;
                     output_tensor_extractor.extract::<TOut>()
                 })
@@ -560,6 +548,24 @@ impl<'a> Session<'a> {
     }
 }
 
+unsafe fn get_tensor_dimensions(
+    tensor_info_ptr: *const sys::OrtTensorTypeAndShapeInfo,
+) -> Result<Vec<i64>> {
+    let mut num_dims = 0;
+    let status = g_ort().GetDimensionsCount.unwrap()(tensor_info_ptr, &mut num_dims);
+    status_to_result(status).map_err(OrtError::GetDimensionsCount)?;
+    assert_ne!(num_dims, 0);
+
+    let mut node_dims: Vec<i64> = vec![0; num_dims as usize];
+    let status = g_ort().GetDimensions.unwrap()(
+        tensor_info_ptr,
+        node_dims.as_mut_ptr(), // FIXME: UB?
+        num_dims,
+    );
+    status_to_result(status).map_err(OrtError::GetDimensions)?;
+    Ok(node_dims)
+}
+
 /// This module contains dangerous functions working on raw pointers.
 /// Those functions are only to be used from inside the
 /// `SessionBuilder::with_model_from_file()` method.
@@ -694,22 +700,7 @@ mod dangerous {
 
         // info!("{} : type={}", i, type_);
 
-        // print input shapes/dims
-        let mut num_dims = 0;
-        let status = unsafe { g_ort().GetDimensionsCount.unwrap()(tensor_info_ptr, &mut num_dims) };
-        status_to_result(status).map_err(OrtError::GetDimensionsCount)?;
-        assert_ne!(num_dims, 0);
-
-        // info!("{} : num_dims={}", i, num_dims);
-        let mut node_dims: Vec<i64> = vec![0; num_dims as usize];
-        let status = unsafe {
-            g_ort().GetDimensions.unwrap()(
-                tensor_info_ptr,
-                node_dims.as_mut_ptr(), // FIXME: UB?
-                num_dims,
-            )
-        };
-        status_to_result(status).map_err(OrtError::GetDimensions)?;
+        let node_dims = unsafe { get_tensor_dimensions(tensor_info_ptr)? };
 
         // for j in 0..num_dims {
         //     info!("{} : dim {}={}", i, j, node_dims[j as usize]);

onnxruntime/tests/data/upsample.onnx

@@ -204,6 +204,94 @@ mod download {
             IMAGE_TO_LOAD, probabilities[0].0
         );
     }
+
+    // This test verifies that dynamically sized inputs and outputs work. It loads and runs
+    // upsample.onnx, which was produced via:
+    //
+    // ```
+    // import subprocess
+    // from tensorflow import keras
+    //
+    // m = keras.Sequential([
+    //     keras.layers.UpSampling2D(size=2)
+    // ])
+    // m.build(input_shape=(None, None, None, 3))
+    // m.summary()
+    // m.save('saved_model')
+    //
+    // subprocess.check_call([
+    //     'python', '-m', 'tf2onnx.convert',
+    //     '--saved-model', 'saved_model',
+    //     '--opset', '12',
+    //     '--output', 'upsample.onnx',
+    // ])
+    // ```
+    #[test]
+    fn upsample() {
+        const IMAGE_TO_LOAD: &str = "mushroom.png";
+
+        let environment = Environment::builder()
+            .with_name("integration_test")
+            .with_log_level(LoggingLevel::Warning)
+            .build()
+            .unwrap();
+
+        let mut session = environment
+            .new_session_builder()
+            .unwrap()
+            .with_optimization_level(GraphOptimizationLevel::Basic)
+            .unwrap()
+            .with_number_threads(1)
+            .unwrap()
+            .with_model_from_file(
+                Path::new(env!("CARGO_MANIFEST_DIR"))
+                    .join("tests")
+                    .join("data")
+                    .join("upsample.onnx"),
+            )
+            .expect("Could not open model from file");
+
+        assert_eq!(
+            session.inputs[0].dimensions().collect::<Vec<_>>(),
+            [None, None, None, Some(3)]
+        );
+        assert_eq!(
+            session.outputs[0].dimensions().collect::<Vec<_>>(),
+            [None, None, None, Some(3)]
+        );
+
+        // Load image, converting to RGB format
+        let image_buffer: ImageBuffer<Rgb<u8>, Vec<u8>> = image::open(
+            Path::new(env!("CARGO_MANIFEST_DIR"))
+                .join("tests")
+                .join("data")
+                .join(IMAGE_TO_LOAD),
+        )
+        .unwrap()
+        .to_rgb();
+
+        let array = ndarray::Array::from_shape_fn((1, 224, 224, 3), |(_, j, i, c)| {
+            let pixel = image_buffer.get_pixel(i as u32, j as u32);
+            let channels = pixel.channels();
+
+            // range [0, 255] -> range [0, 1]
+            (channels[c] as f32) / 255.0
+        });
+
+        // Just one input
+        let input_tensor_values = vec![array];
+
+        // Perform the inference
+        let outputs: Vec<
+            onnxruntime::tensor::OrtOwnedTensor<f32, ndarray::Dim<ndarray::IxDynImpl>>,
+        > = session.run(input_tensor_values).unwrap();
+
+        assert_eq!(outputs.len(), 1);
+        let output = &outputs[0];
+
+        // The image should have doubled in size
+        assert_eq!(output.shape(), [1, 448, 448, 3]);
+    }
 }
 
 fn get_imagenet_labels() -> Result<Vec<String>, io::Error> {