Merge pull request #647 from robertknight/pytorch-half-pixel

Support `pytorch_half_pixel` value for `coordinate_transform_mode`

Author: robertknight

rten-convert/rten_convert/schema_generated.py

@@ -143,6 +143,7 @@ class CoordTransformMode(object):
     HalfPixel = 0
     Asymmetric = 1
     AlignCorners = 2
+    PytorchHalfPixel = 3
 
 
 class NearestMode(object):

src/model_builder.rs

@@ -796,6 +796,9 @@ impl<'mb, 'a> GraphBuilder<'mb, 'a> {
                     CoordTransformMode::Asymmetric => sg::CoordTransformMode::Asymmetric,
                     CoordTransformMode::HalfPixel => sg::CoordTransformMode::HalfPixel,
                     CoordTransformMode::AlignCorners => sg::CoordTransformMode::AlignCorners,
+                    CoordTransformMode::PytorchHalfPixel => {
+                        sg::CoordTransformMode::PytorchHalfPixel
+                    }
                 };
                 let nearest_mode = match args.nearest_mode {
                     NearestMode::Ceil => sg::NearestMode::Ceil,

src/op_registry.rs

@@ -817,6 +817,7 @@ impl_read_op!(Resize, attrs_as_resize_attrs, |attrs: sg::ResizeAttrs| {
         sg::CoordTransformMode::Asymmetric => CoordTransformMode::Asymmetric,
         sg::CoordTransformMode::HalfPixel => CoordTransformMode::HalfPixel,
         sg::CoordTransformMode::AlignCorners => CoordTransformMode::AlignCorners,
+        sg::CoordTransformMode::PytorchHalfPixel => CoordTransformMode::PytorchHalfPixel,
         _ => CoordTransformMode::default(),
     };
 

src/ops/resize.rs

@@ -51,6 +51,18 @@ fn input_coord(
         Ctm::AlignCorners => {
             dest_coord as f32 * (length_original - 1) as f32 / (length_resized - 1) as f32
         }
+        Ctm::PytorchHalfPixel => {
+            // There are some queries over this transform mode, see
+            // https://github.com/onnx/onnx/issues/4275 (applies to cubic interpolation only)
+            // and https://github.com/onnx/onnx/issues/4276 (comparison with
+            // PyTorch behavior). This implementation does however match
+            // ONNX Runtime (https://github.com/microsoft/onnxruntime/blob/24620e70d9f14956a0dc84bb8a332dcd64c95a94/onnxruntime/core/providers/cpu/tensor/upsamplebase.h#L331)
+            if length_resized > 1 {
+                (dest_coord as f32 + 0.5) / scale - 0.5
+            } else {
+                0.
+            }
+        }
     }
 }
 
@@ -73,6 +85,7 @@ pub enum CoordTransformMode {
     HalfPixel,
     Asymmetric,
     AlignCorners,
+    PytorchHalfPixel,
 }
 
 const CHAN_GROUP_SIZE: usize = 4;
@@ -648,18 +661,25 @@ mod tests {
                 coord_transform_mode: None,
                 expected: Tensor::from_data(&[1, 1, 0, 0], vec![]),
             },
-            // Scale width and height by 0.5x
+            // Scale to output width and height less than 2, using `HalfPixel`
+            // `coord_transform_mode`.
+            //
+            // When the output size is < 2, `half_pixel` and `pytorch_half_pixel`
+            // produce different results. Otherwise they are the same.
             Case {
                 image,
                 scales: vec![1., 1., 0.5, 0.5],
-                coord_transform_mode: None,
-
-                // OpenCV and PyTorch produce different results for this case.
-                // This result matches OpenCV. This relates to the `half_pixel`
-                // vs `pytorch_half_pixel` values for the `coordinate_transformation_mode`
-                // attribute in the ONNX op.
+                coord_transform_mode: Some(CoordTransformMode::HalfPixel),
                 expected: Tensor::from_data(&[1, 1, 1, 1], vec![0.5]),
             },
+            // Scale to output width and height less than 2, using `PytorchHalfPixel`
+            // `coord_transform_mode`.
+            Case {
+                image,
+                scales: vec![1., 1., 0.5, 0.5],
+                coord_transform_mode: Some(CoordTransformMode::PytorchHalfPixel),
+                expected: Tensor::from_data(&[1, 1, 1, 1], vec![0.2]),
+            },
             // Scale width and height by 1x
             Case {
                 image,

src/schema.fbs

@@ -155,7 +155,8 @@ enum DataType: ubyte {
 enum CoordTransformMode: ubyte {
   HalfPixel,
   Asymmetric,
-  AlignCorners
+  AlignCorners,
+  PytorchHalfPixel,
 }
 
 // Rounding modes supported by Resize operator when `ResizeMode` is `Nearest`.

src/schema_generated.rs

@@ -835,16 +835,17 @@ pub const ENUM_MIN_COORD_TRANSFORM_MODE: u8 = 0;
     since = "2.0.0",
     note = "Use associated constants instead. This will no longer be generated in 2021."
 )]
-pub const ENUM_MAX_COORD_TRANSFORM_MODE: u8 = 2;
+pub const ENUM_MAX_COORD_TRANSFORM_MODE: u8 = 3;
 #[deprecated(
     since = "2.0.0",
     note = "Use associated constants instead. This will no longer be generated in 2021."
 )]
 #[allow(non_camel_case_types)]
-pub const ENUM_VALUES_COORD_TRANSFORM_MODE: [CoordTransformMode; 3] = [
+pub const ENUM_VALUES_COORD_TRANSFORM_MODE: [CoordTransformMode; 4] = [
     CoordTransformMode::HalfPixel,
     CoordTransformMode::Asymmetric,
     CoordTransformMode::AlignCorners,
+    CoordTransformMode::PytorchHalfPixel,
 ];
 
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
@@ -855,17 +856,23 @@ impl CoordTransformMode {
     pub const HalfPixel: Self = Self(0);
     pub const Asymmetric: Self = Self(1);
     pub const AlignCorners: Self = Self(2);
+    pub const PytorchHalfPixel: Self = Self(3);
 
     pub const ENUM_MIN: u8 = 0;
-    pub const ENUM_MAX: u8 = 2;
-    pub const ENUM_VALUES: &'static [Self] =
-        &[Self::HalfPixel, Self::Asymmetric, Self::AlignCorners];
+    pub const ENUM_MAX: u8 = 3;
+    pub const ENUM_VALUES: &'static [Self] = &[
+        Self::HalfPixel,
+        Self::Asymmetric,
+        Self::AlignCorners,
+        Self::PytorchHalfPixel,
+    ];
     /// Returns the variant's name or "" if unknown.
     pub fn variant_name(self) -> Option<&'static str> {
         match self {
             Self::HalfPixel => Some("HalfPixel"),
             Self::Asymmetric => Some("Asymmetric"),
             Self::AlignCorners => Some("AlignCorners"),
+            Self::PytorchHalfPixel => Some("PytorchHalfPixel"),
             _ => None,
         }
     }