feat: add support for SpaceToDepth onnx node

Author: tye-singwa

crates/burn-import/onnx-tests/build.rs

@@ -122,6 +122,7 @@ fn main() {
         .input("tests/slice/slice.onnx")
         .input("tests/slice/slice_shape.onnx")
         .input("tests/softmax/softmax.onnx")
+        .input("tests/space_to_depth/space_to_depth.onnx")
         .input("tests/sqrt/sqrt.onnx")
         .input("tests/squeeze/squeeze_multiple.onnx")
         .input("tests/squeeze/squeeze.onnx")

crates/burn-import/onnx-tests/tests/space_to_depth/mod.rs

@@ -0,0 +1,55 @@
+// Import the shared macro
+use crate::include_models;
+include_models!(space_to_depth);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Tensor, TensorData, Tolerance, ops::FloatElem};
+
+    type Backend = burn_ndarray::NdArray<f32>;
+    type FT = FloatElem<Backend>;
+
+    #[test]
+    fn space_to_depth() {
+        let device = Default::default();
+        let model: space_to_depth::Model<Backend> = space_to_depth::Model::new(&device);
+
+        let input = Tensor::<Backend, 4>::from_floats(
+            [
+                [[
+                    [0.5, -0.14, 0.65, 1.52, -0.23, -0.23],
+                    [1.58, 0.77, -0.47, 0.54, -0.46, -0.47],
+                    [0.24, -1.91, -1.72, -0.56, -1.01, 0.31],
+                    [-0.91, -1.41, 1.47, -0.23, 0.07, -1.42],
+                ]],
+                [[
+                    [-0.54, 0.11, -1.15, 0.38, -0.6, -0.29],
+                    [-0.6, 1.85, -0.01, -1.06, 0.82, -1.22],
+                    [0.21, -1.96, -1.33, 0.2, 0.74, 0.17],
+                    [-0.12, -0.3, -1.48, -0.72, -0.46, 1.06],
+                ]],
+            ],
+            &device,
+        );
+        let output = model.forward(input);
+        let expected = TensorData::from([
+            [
+                [[0.5, 0.65, -0.23], [0.24, -1.72, -1.01]],
+                [[-0.14, 1.52, -0.23], [-1.91, -0.56, 0.31]],
+                [[1.58, -0.47, -0.46], [-0.91, 1.47, 0.07]],
+                [[0.77, 0.54, -0.47], [-1.41, -0.23, -1.42]],
+            ],
+            [
+                [[-0.54, -1.15, -0.6], [0.21, -1.33, 0.74]],
+                [[0.11, 0.38, -0.29], [-1.96, 0.2, 0.17]],
+                [[-0.6, -0.01, 0.82], [-0.12, -1.48, -0.46]],
+                [[1.85, -1.06, -1.22], [-0.3, -0.72, 1.06]],
+            ],
+        ]);
+
+        output
+            .to_data()
+            .assert_approx_eq::<FT>(&expected, Tolerance::default());
+    }
+}

crates/burn-import/onnx-tests/tests/space_to_depth/space_to_depth.onnx

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+
+# used to generate models: space_to_depth_*.onnx
+
+import numpy as np
+import onnx
+import onnx.helper
+from onnx import TensorProto
+from onnx.reference import ReferenceEvaluator
+
+
+def build_model():
+    # Define the graph inputs and outputs
+    input = onnx.helper.make_tensor_value_info(
+        'input', TensorProto.FLOAT, [2, 1, 4, 6])
+    output = onnx.helper.make_tensor_value_info(
+        'output', TensorProto.FLOAT, [2, 4, 2, 3])
+
+    # Create the SpaceToDepth node
+    space_to_depth = onnx.helper.make_node(
+        "SpaceToDepth",
+        inputs=["input"],
+        outputs=["output"],
+        name="SpaceToDepthNode",
+        blocksize=2,
+    )
+
+    # Create the graph
+    graph = onnx.helper.make_graph(
+        [space_to_depth],
+        'SpaceToDepthModel',
+        [input],
+        [output],
+    )
+
+    # Create the model
+    model = onnx.helper.make_model(
+        opset_imports=[onnx.helper.make_operatorsetid("", 21)],
+        graph=graph,
+        producer_name='ONNX_Generator',
+    )
+
+    return model
+
+
+if __name__ == "__main__":
+    # Set seed and precision
+    np.random.seed(42)
+    np.set_printoptions(precision=8)
+
+    # Build model
+    test_input = np.random.randn(2, 1, 4, 6).round(2)
+    onnx_model = build_model()
+    file_name = "space_to_depth.onnx"
+
+    # Ensure valid ONNX and save
+    onnx.checker.check_model(onnx_model)
+    onnx.save(onnx_model, file_name)
+    print(f"Finished exporting model to {file_name}")
+
+    # Output some test data for use in the test
+    print(f"Test input data:\n{repr(test_input)}")
+    print(f"Test input data shape: {test_input.shape}")
+    session = ReferenceEvaluator(file_name, verbose=1)
+    test_output, = session.run(None, {"input": test_input})
+    print(f"Test output:\n{repr(test_output)}")
+    print(f"Test output shape: {test_output.shape}")

crates/burn-import/onnx-tests/tests/space_to_depth/space_to_depth.py

@@ -18,12 +18,12 @@ pub mod conv;
 pub mod conv_transpose;
 pub mod cos;
 pub mod cosh;
+pub mod depth_to_space;
 pub mod div;
 pub mod dropout;
 pub mod equal;
 pub mod erf;
 pub mod exp;
-pub mod depth_to_space;
 pub mod expand;
 pub mod flatten;
 pub mod floor;
@@ -80,6 +80,7 @@ pub mod sin;
 pub mod sinh;
 pub mod slice;
 pub mod softmax;
+pub mod space_to_depth;
 pub mod split;
 pub mod sqrt;
 pub mod squeeze;

crates/burn-import/onnx-tests/tests/test_mod.rs

@@ -6,8 +6,8 @@ use super::{
     concat::ConcatNode, constant::ConstantNode, constant_of_shape::ConstantOfShapeNode,
     conv_transpose_1d::ConvTranspose1dNode, conv_transpose_2d::ConvTranspose2dNode,
     conv_transpose_3d::ConvTranspose3dNode, conv1d::Conv1dNode, conv2d::Conv2dNode,
-    conv3d::Conv3dNode, dropout::DropoutNode, expand::ExpandNode, floor::FloorNode,
-    gather::GatherNode, gather_elements::GatherElementsNode, gemm::GemmNode,
+    conv3d::Conv3dNode, depth_to_space::DepthToSpaceNode, dropout::DropoutNode, expand::ExpandNode,
+    floor::FloorNode, gather::GatherNode, gather_elements::GatherElementsNode, gemm::GemmNode,
     global_avg_pool::GlobalAvgPoolNode, group_norm::GroupNormNode, instance_norm::InstanceNormNode,
     layer_norm::LayerNormNode, linear::LinearNode, mask_where::WhereNode, matmul::MatmulNode,
     max_pool1d::MaxPool1dNode, max_pool2d::MaxPool2dNode, mean::MeanNode, one_hot::OneHotNode,
@@ -17,9 +17,8 @@ use super::{
     resize::ResizeNode, round::RoundNode, slice::SliceNode, split::SplitNode, squeeze::SqueezeNode,
     sum::SumNode, tile::TileNode, top_k::TopKNode, trilu::TriluNode, unary::UnaryNode,
     unsqueeze::UnsqueezeNode,
-    depth_to_space::DepthToSpaceNode
 };
-use crate::burn::{BurnImports, Scope, Type};
+use crate::burn::{BurnImports, Scope, Type, node::space_to_depth::SpaceToDepthNode};
 use burn::record::PrecisionSettings;
 use proc_macro2::TokenStream;
 use serde::Serialize;
@@ -126,6 +125,7 @@ pub enum Node<PS: PrecisionSettings> {
     Round(RoundNode),
     Slice(SliceNode),
     Squeeze(SqueezeNode),
+    SpaceToDepth(SpaceToDepthNode),
     Split(SplitNode),
     Sum(SumNode),
     Tile(TileNode),
@@ -187,6 +187,7 @@ macro_rules! match_all {
             Node::Resize(node) => $func(node),
             Node::Round(node) => $func(node),
             Node::Slice(node) => $func(node),
+            Node::SpaceToDepth(node) => $func(node),
             Node::Squeeze(node) => $func(node),
             Node::Sum(node) => $func(node),
             Node::Tile(node) => $func(node),
@@ -257,6 +258,7 @@ impl<PS: PrecisionSettings> Node<PS> {
             Node::Resize(_) => "resize",
             Node::Round(_) => "round",
             Node::Slice(_) => "slice",
+            Node::SpaceToDepth(_) => "space_to_depth",
             Node::Squeeze(_) => "squeeze",
             Node::Sum(_) => "add",
             Node::Tile(_) => "tile",
@@ -399,10 +401,9 @@ pub(crate) mod tests {
             ),
         ));
 
-        graph.register_input_output(
-            vec!["tensor1".to_string(), "tensor2".to_string()],
-            vec!["tensor4".to_string()],
-        );
+        graph.register_input_output(vec!["tensor1".to_string(), "tensor2".to_string()], vec![
+            "tensor4".to_string(),
+        ]);
 
         let expected = quote! {
             use burn::{
@@ -485,10 +486,9 @@ pub(crate) mod tests {
             TensorType::new_float("output", 4),
         ));
 
-        graph.register_input_output(
-            vec!["tensor1".to_string(), "tensor2".to_string()],
-            vec!["output".to_string()],
-        );
+        graph.register_input_output(vec!["tensor1".to_string(), "tensor2".to_string()], vec![
+            "output".to_string(),
+        ]);
 
         let expected = quote! {
             use burn::{

crates/burn-import/src/burn/node/base.rs

@@ -80,7 +80,7 @@ mod tests {
 
         graph.register(DepthToSpaceNode::new(
             TensorType::new_float("input", 4),
-            TensorType::new_float("output", 5),
+            TensorType::new_float("output", 4),
             DepthToSpaceConfig::new(DepthToSpaceMode::DCR, 2),
         ));
 
@@ -105,7 +105,7 @@ mod tests {
                     }
                 }
                 #[allow(clippy::let_and_return, clippy::approx_constant)]
-                pub fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 5> {
+                pub fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
                     let output = {
                         let [b, c, h, w] = input.shape().dims();
                         input
@@ -152,7 +152,7 @@ mod tests {
                     }
                 }
                 #[allow(clippy::let_and_return, clippy::approx_constant)]
-                pub fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 5> {
+                pub fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
                     let output = {
                         let [b, c, h, w] = input.shape().dims();
                         input

crates/burn-import/src/burn/node/depth_to_space.rs

@@ -45,6 +45,7 @@ pub(crate) mod reshape;
 pub(crate) mod resize;
 pub(crate) mod round;
 pub(crate) mod slice;
+pub(crate) mod space_to_depth;
 pub(crate) mod split;
 pub(crate) mod squeeze;
 pub(crate) mod sum;

crates/burn-import/src/burn/node/mod.rs

@@ -0,0 +1,106 @@
+use super::{Node, NodeCodegen};
+use crate::burn::{Scope, TensorType, Type};
+use burn::record::PrecisionSettings;
+use proc_macro2::TokenStream;
+use quote::quote;
+
+#[derive(Debug, Clone)]
+pub struct SpaceToDepthNode {
+    pub input: TensorType,
+    pub output: TensorType,
+    pub block_size: usize,
+}
+
+impl SpaceToDepthNode {
+    pub fn new(input: TensorType, output: TensorType, block_size: usize) -> Self {
+        Self {
+            input,
+            output,
+            block_size,
+        }
+    }
+}
+
+impl<PS: PrecisionSettings> NodeCodegen<PS> for SpaceToDepthNode {
+    fn input_types(&self) -> Vec<Type> {
+        vec![Type::Tensor(self.input.clone())]
+    }
+
+    fn output_types(&self) -> Vec<Type> {
+        vec![Type::Tensor(self.output.clone())]
+    }
+
+    fn forward(&self, scope: &mut Scope, node_position: usize) -> TokenStream {
+        let input = scope.tensor_use_owned(&self.input, node_position);
+        let output = &self.output.name;
+        let block_size = self.block_size;
+
+        quote! {
+            let #output = {
+                let [b, c, h, w] = #input.shape().dims();
+                #input
+                    .reshape([b, c, h / #block_size, #block_size, w / #block_size, #block_size])
+                    .permute([0, 3, 5, 1, 2, 4])
+                    .reshape([b, c * #block_size * #block_size, h / #block_size, w / #block_size])
+            };
+        }
+    }
+
+    fn into_node(self) -> Node<PS> {
+        Node::SpaceToDepth(self)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::burn::{TensorType, graph::BurnGraph, node::test::assert_tokens};
+    use burn::record::FullPrecisionSettings;
+
+    #[test]
+    fn test_codegen() {
+        let mut graph = BurnGraph::<FullPrecisionSettings>::default();
+
+        graph.register(SpaceToDepthNode::new(
+            TensorType::new_float("input", 4),
+            TensorType::new_float("output", 4),
+            2,
+        ));
+
+        graph.register_input_output(vec!["input".to_string()], vec!["output".to_string()]);
+
+        let expected = quote! {
+            use burn::{
+                module::Module,
+                tensor::{backend::Backend, Tensor},
+            };
+            #[derive(Module, Debug)]
+            pub struct Model<B: Backend> {
+                phantom: core::marker::PhantomData<B>,
+                device: burn::module::Ignored<B::Device>,
+            }
+            impl<B: Backend> Model<B> {
+                #[allow(unused_variables)]
+                pub fn new(device: &B::Device) -> Self {
+                    Self {
+                        phantom: core::marker::PhantomData,
+                        device: burn::module::Ignored(device.clone()),
+                    }
+                }
+                #[allow(clippy::let_and_return, clippy::approx_constant)]
+                pub fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
+                    let output = {
+                        let [b, c, h, w] = input.shape().dims();
+                        input
+                            .reshape([b, c, h / 2usize, 2usize, w / 2usize, 2usize])
+                            .permute([0, 3, 5, 1, 2, 4])
+                            .reshape([b, c * 2usize * 2usize, h / 2usize, w / 2usize])
+                    };
+                    output
+                }
+            }
+        };
+
+        assert_tokens(graph.codegen(), expected);
+    }
+}