Changes and fixes required for squeezenet1.1-opset7 onnx (#639)

Author: antimora

burn-core/src/nn/padding.rs

@@ -34,7 +34,7 @@ impl PaddingConfig1d {
 }
 
 /// Padding configuration for 2D operators.
-#[derive(Module, Config, Debug)]
+#[derive(Module, Config, Debug, PartialEq)]
 pub enum PaddingConfig2d {
     /// Dynamically calculate the amount of padding necessary to ensure that the output size will be
     /// the same as the input.

burn-core/src/nn/pool/avg_pool1d.rs

@@ -21,6 +21,20 @@ pub struct AvgPool1dConfig {
 }
 
 /// Applies a 1D avg pooling over input tensors.
+///
+/// See [AvgPool1dConfig](AvgPool1dConfig) for details.
+///
+/// # Remarks
+///
+/// The zero-padding values will be included in the calculation
+/// of the average. This means that the zeros are counted as
+/// legitimate values, and they contribute to the denominator
+/// when calculating the average. This is equivalent to
+/// `torch.nn.AvgPool2d` with `count_include_pad=True`.
+///
+/// TODO: Add support for `count_include_pad=False`, see
+/// [Issue 636](https://github.com/burn-rs/burn/issues/636)
+
 #[derive(Module, Debug, Clone)]
 pub struct AvgPool1d {
     stride: usize,

burn-core/src/nn/pool/avg_pool2d.rs

@@ -8,7 +8,7 @@ use crate::tensor::Tensor;
 use burn_tensor::module::avg_pool2d;
 
 /// Configuration to create a [2D avg pooling](AvgPool2d) layer.
-#[derive(Config)]
+#[derive(Config, Debug)]
 pub struct AvgPool2dConfig {
     /// The size of the kernel.
     pub kernel_size: [usize; 2],
@@ -21,6 +21,19 @@ pub struct AvgPool2dConfig {
 }
 
 /// Applies a 2D avg pooling over input tensors.
+///
+/// See [AvgPool2dConfig](AvgPool2dConfig) for details.
+///
+/// # Remarks
+///
+/// The zero-padding values will be included in the calculation
+/// of the average. This means that the zeros are counted as
+/// legitimate values, and they contribute to the denominator
+/// when calculating the average. This is equivalent to
+/// `torch.nn.AvgPool2d` with `count_include_pad=True`.
+///
+/// TODO: Add support for `count_include_pad=False`, see
+/// [Issue 636](https://github.com/burn-rs/burn/issues/636)
 #[derive(Module, Debug, Clone)]
 pub struct AvgPool2d {
     stride: [usize; 2],

burn-import/README.md

@@ -23,7 +23,8 @@ List taken from [here](https://github.com/onnx/onnx/blob/main/docs/Operators.md)
 - [ ] Asinh
 - [ ] Atan
 - [ ] Atanh
-- [ ] AveragePool
+- [ ] AveragePool1d
+- [x] AveragePool2d
 - [x] BatchNormalization
 - [ ] Bernoulli
 - [ ] BitShift
@@ -107,7 +108,6 @@ List taken from [here](https://github.com/onnx/onnx/blob/main/docs/Operators.md)
 - [ ] MatMul
 - [ ] MatMulInteger
 - [ ] Max
-- [ ] MaxPool
 - [ ] MaxPool1d
 - [x] MaxPool2d
 - [ ] MaxRoiPool

burn-import/onnx-tests/build.rs

@@ -7,17 +7,20 @@ fn main() {
     // Add onnx models.
     ModelGen::new()
         .input("tests/add/add.onnx")
-        .input("tests/sub/sub.onnx")
-        .input("tests/mul/mul.onnx")
-        .input("tests/div/div.onnx")
+        .input("tests/avg_pool2d/avg_pool2d.onnx")
         .input("tests/concat/concat.onnx")
         .input("tests/conv1d/conv1d.onnx")
         .input("tests/conv2d/conv2d.onnx")
-        .input("tests/dropout/dropout.onnx")
+        .input("tests/div/div.onnx")
+        .input("tests/dropout/dropout_opset16.onnx")
+        .input("tests/dropout/dropout_opset7.onnx")
         .input("tests/global_avr_pool/global_avr_pool.onnx")
-        .input("tests/softmax/softmax.onnx")
         .input("tests/log_softmax/log_softmax.onnx")
         .input("tests/maxpool2d/maxpool2d.onnx")
+        .input("tests/mul/mul.onnx")
+        .input("tests/reshape/reshape.onnx")
+        .input("tests/softmax/softmax.onnx")
+        .input("tests/sub/sub.onnx")
         .out_dir("model/")
         .run_from_script();
 

burn-import/onnx-tests/tests/avg_pool2d/avg_pool2d.onnx

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+
+# used to generate model: avg_pool2d.onnx
+
+import torch
+import torch.nn as nn
+
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        # TODO when https://github.com/burn-rs/burn/issues/636 is resolved, test this with a model 
+        # that uses `count_include_pad=False` and padding=(2, 1) 
+        self.pool2d = nn.AvgPool2d((4, 2), stride=(
+            2, 1), padding=(0, 0), count_include_pad=False)
+
+    def forward(self, x):
+        x = self.pool2d(x)
+        return x
+
+
+def main():
+    # Set seed for reproducibility
+    torch.manual_seed(3)
+
+    # Print options
+    torch.set_printoptions(precision=3)
+
+    # Export to onnx
+    model = Model()
+    model.eval()
+    device = torch.device("cpu")
+
+    file_name = "avg_pool2d.onnx"
+    test_input = torch.randn(1, 1, 5, 5, device=device)
+    torch.onnx.export(model, test_input, file_name,
+                      verbose=False, opset_version=16)
+
+    print("Finished exporting model to {}".format(file_name))
+
+    # Output some test data for use in the test
+    print("Test input data shape of ones: {}".format(test_input.shape))
+    print("Test input data of ones: {}".format(test_input))
+    output = model.forward(test_input)
+    print("Test output data shape: {}".format(output.shape))
+    print("Test output: {}".format(output))
+
+
+if __name__ == '__main__':
+    main()

burn-import/onnx-tests/tests/avg_pool2d/avg_pool2d.py

@@ -23,7 +23,7 @@ def main():
     model.eval()
     device = torch.device("cpu")
 
-    file_name = "dropout.onnx"
+    file_name = "dropout_opset16.onnx"
     test_input = torch.ones(2, 4, 10, 15, device=device)
     torch.onnx.export(model, test_input, file_name,
                       training=torch.onnx.TrainingMode.TRAINING,

burn-import/onnx-tests/tests/dropout/dropout.py

@@ -12,17 +12,20 @@ macro_rules! include_models {
 // ATTENTION: Modify this macro to include all models in the `model` directory.
 include_models!(
     add,
-    sub,
-    mul,
-    div,
+    avg_pool2d,
     concat,
     conv1d,
     conv2d,
-    dropout,
+    div,
+    dropout_opset16,
+    dropout_opset7,
     global_avr_pool,
-    softmax,
     log_softmax,
-    maxpool2d
+    maxpool2d,
+    mul,
+    reshape,
+    softmax,
+    sub
 );
 
 #[cfg(test)]
@@ -151,8 +154,27 @@ mod tests {
     }
 
     #[test]
-    fn dropout() {
-        let model: dropout::Model<Backend> = dropout::Model::default();
+    fn dropout_opset16() {
+        let model: dropout_opset16::Model<Backend> = dropout_opset16::Model::default();
+
+        // Run the model with ones as input for easier testing
+        let input = Tensor::<Backend, 4>::ones([2, 4, 10, 15]);
+
+        let output = model.forward(input);
+
+        let expected_shape = Shape::from([2, 4, 10, 15]);
+        assert_eq!(output.shape(), expected_shape);
+
+        let output_sum = output.sum().into_scalar();
+
+        let expected_sum = 1200.0; // from pytorch
+
+        assert!(expected_sum.approx_eq(output_sum, (1.0e-4, 2)));
+    }
+
+    #[test]
+    fn dropout_opset7() {
+        let model: dropout_opset7::Model<Backend> = dropout_opset7::Model::default();
 
         // Run the model with ones as input for easier testing
         let input = Tensor::<Backend, 4>::ones([2, 4, 10, 15]);
@@ -255,4 +277,36 @@ mod tests {
 
         assert_eq!(output.to_data(), expected);
     }
+
+    #[test]
+    fn avg_pool2d() {
+        // Initialize the model without weights (because the exported file does not contain them)
+        let model: avg_pool2d::Model<Backend> = avg_pool2d::Model::new();
+
+        // Run the model
+        let input = Tensor::<Backend, 4>::from_floats([[[
+            [-0.077, 0.360, -0.782, 0.072, 0.665],
+            [-0.287, 1.621, -1.597, -0.052, 0.611],
+            [0.760, -0.034, -0.345, 0.494, -0.078],
+            [-1.805, -0.476, 0.205, 0.338, 1.353],
+            [0.374, 0.013, 0.774, -0.109, -0.271],
+        ]]]);
+        let output = model.forward(input);
+        let expected = Data::from([[[[0.008, -0.131, -0.208, 0.425]]]]);
+
+        output.to_data().assert_approx_eq(&expected, 3);
+    }
+
+    #[test]
+    fn reshape() {
+        // Initialize the model without weights (because the exported file does not contain them)
+        let model: reshape::Model<Backend> = reshape::Model::new();
+
+        // Run the model
+        let input = Tensor::<Backend, 1>::from_floats([0., 1., 2., 3.]);
+        let output = model.forward(input);
+        let expected = Data::from([[0., 1., 2., 3.]]);
+
+        assert_eq!(output.to_data(), expected);
+    }
 }

burn-import/onnx-tests/tests/dropout/dropout.onnx burn-import/onnx-tests/tests/dropout/dropout_opset16.onnx

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+
+# used to generate model: reshape.onnx
+
+import torch
+import torch.nn as nn
+
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x):
+        x = x.reshape(2, 2)
+        x = x.reshape(1, -1)  # -1 means infer from other dimensions
+        return x
+
+
+def main():
+
+    # Set seed for reproducibility
+    torch.manual_seed(42)
+
+    torch.set_printoptions(precision=8)
+
+    # Export to onnx
+    model = Model()
+    model.eval()
+    device = torch.device("cpu")
+
+    file_name = "reshape.onnx"
+    test_input = torch.arange(4., device=device)
+    torch.onnx.export(model, test_input, file_name,
+                      verbose=False, opset_version=16)
+
+    print("Finished exporting model to {}".format(file_name))
+
+    # Output some test data for use in the test
+    print("Test input data of ones: {}".format(test_input))
+    print("Test input data shape of ones: {}".format(test_input.shape))
+    output = model.forward(test_input)
+    print("Test output data shape: {}".format(output.shape))
+
+    print("Test output: {}".format(output))
+
+
+if __name__ == '__main__':
+    main()