From 0616ef1d731cb9df8ba5feb7edf3342d0d3dba46 Mon Sep 17 00:00:00 2001
From: hrouis <houssemeddine.rouis92@gmail.com>
Date: Fri, 24 Jan 2025 16:08:52 +0100
Subject: [PATCH] add dilations to maxpooling

---
 include/aidge/operator/MaxPooling.hpp         | 22 +++++++++++++++---
 python_binding/operator/pybind_MaxPooling.cpp | 11 ++++++++-
 src/operator/MaxPooling.cpp                   | 23 ++++++++++++-------
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/include/aidge/operator/MaxPooling.hpp b/include/aidge/operator/MaxPooling.hpp
index 8503b1be1..4b1190c14 100644
--- a/include/aidge/operator/MaxPooling.hpp
+++ b/include/aidge/operator/MaxPooling.hpp
@@ -41,7 +41,10 @@ enum class MaxPoolingAttr {
    * Must be positive integers.
    */
   StrideDims,
-
+  /**
+   * @brief Dilation along each spatial axis. Default value is 1.
+   */
+  Dilations,
   /**
    * @brief Kernel dimensions specifying the size of the pooling window for each spatial dimension.
    * For example, common kernel dimensions include 2x2 or 3x3.
@@ -91,6 +94,7 @@ public:
     static const std::string Type; ///< Static identifier for this operator type.
 
     using Attributes_ = StaticAttributes<MaxPoolingAttr,
+                                         std::array<DimSize_t, DIM>,
                                          std::array<DimSize_t, DIM>,
                                          std::array<DimSize_t, DIM>,
                                          bool>;
@@ -107,10 +111,12 @@ public:
      * @brief Constructor.
      * @param[in] kernel_dims Size of the pooling window for each spatial dimension.
      * @param[in] stride_dims Step size (stride) for sliding the pooling window across input dimensions.
+     * @param[in] dilations Spatial dilations for the pooling operation.
      * @param[in] ceil_mode Indicates whether to use ceil or floor for output size calculation.
      */
     MaxPooling_Op(const std::array<DimSize_t, DIM> &kernel_dims,
                   const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                  const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t, DIM>(1),
                   bool ceil_mode = false);
 
     /**
@@ -159,6 +165,12 @@ public:
      */
     inline std::array<DimSize_t, DIM>& strideDims() const { return mAttributes->template getAttr<MaxPoolingAttr::StrideDims>(); }
 
+    /**
+     * @brief Accessor for dilations.
+     * @return An array representing spatial dilations.
+     */
+    inline std::array<DimSize_t, DIM>& dilations() const { return mAttributes->template getAttr<MaxPoolingAttr::Dilations>(); }
+
     /**
      * @brief Accessor for kernel dimensions.
      * @return An array representing kernel dimensions.
@@ -197,6 +209,7 @@ extern template class Aidge::MaxPooling_Op<3>;
  * @param[in] kernel_dims Kernel dimensions specifying the size of the pooling window.
  * @param[in] name Optional name for the operation.
  * @param[in] stride_dims Stride dimensions specifying the step size for the pooling window.
+ * @param[in] dilations Spatial dilations for the pooling operation.
  * @param[in] ceil_mode Indicates whether to use ceil mode for output size calculation.
  * @return A shared pointer to a Node representing the MaxPooling operation.
  */
@@ -204,6 +217,7 @@ template <std::array<DimSize_t, 1>::size_type DIM>
 std::shared_ptr<Node> MaxPooling(const std::array<DimSize_t, DIM> &kernel_dims,
                                  const std::string& name = "",
                                  const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+                                 const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
                                  bool ceil_mode=false);
 
 /**
@@ -212,6 +226,7 @@ std::shared_ptr<Node> MaxPooling(const std::array<DimSize_t, DIM> &kernel_dims,
  * @param[in] kernel_dims C-style array of kernel dimensions.
  * @param[in] name Optional name for the operation.
  * @param[in] stride_dims Stride dimensions specifying the step size for the pooling window.
+ * @param[in] dilations Spatial dilations for the pooling operation.
  * @param[in] ceil_mode Indicates whether to use ceil mode for output size calculation.
  * @return A shared pointer to a Node representing the MaxPooling operation.
  */
@@ -220,9 +235,10 @@ inline std::shared_ptr<Node> MaxPooling(
     DimSize_t const (&kernel_dims)[DIM],
     const std::string& name = "",
     const std::array<DimSize_t, DIM> &stride_dims = create_array<DimSize_t,DIM>(1),
+    const std::array<DimSize_t, DIM> &dilations = create_array<DimSize_t,DIM>(1),
     bool ceil_mode = false) {
     static_assert(DIM<=MaxDim,"Too many kernel dimensions required by MaxPooling, not supported");
-    return MaxPooling(to_array(kernel_dims), name, stride_dims, ceil_mode);
+    return MaxPooling(to_array(kernel_dims), name, stride_dims, dilations, ceil_mode);
 }
 
 }  // namespace Aidge
@@ -232,7 +248,7 @@ namespace {
  * @brief String representations of MaxPooling attributes for debugging and logging.
  */
 template <>
-const char *const EnumStrings<Aidge::MaxPoolingAttr>::data[] = {"stride_dims", "kernel_dims", "ceil_mode"};
+const char *const EnumStrings<Aidge::MaxPoolingAttr>::data[] = {"stride_dims", "kernel_dims", "dilations", "ceil_mode"};
 }
 
 #endif /* AIDGE_CORE_OPERATOR_MAXPOOLING_H_ */
diff --git a/python_binding/operator/pybind_MaxPooling.cpp b/python_binding/operator/pybind_MaxPooling.cpp
index 8834625a8..bdbc1edd3 100644
--- a/python_binding/operator/pybind_MaxPooling.cpp
+++ b/python_binding/operator/pybind_MaxPooling.cpp
@@ -37,14 +37,18 @@ template <DimIdx_t DIM> void declare_MaxPoolingOp(py::module &m) {
         :type kernel_dims: List[int]
         :param stride_dims: The stride (step size) to move the kernel over the input.
         :type stride_dims: List[int]
+        :param dilations: The dilation value along each spatial axis of filter.
+        :type dilations: List[int]
         :param ceil_mode: Whether to use ceil or floor when calculating the output dimensions.
         :type ceil_mode: bool
     )mydelimiter")
   .def(py::init<const std::array<DimSize_t, DIM> &,
+                const std::array<DimSize_t, DIM> &,
                 const std::array<DimSize_t, DIM> &,
                 bool>(),
         py::arg("kernel_dims"),
         py::arg("stride_dims"),
+        py::arg("dilations"),
         py::arg("ceil_mode"))
   .def_static("get_inputs_name", &MaxPooling_Op<DIM>::getInputsName)
   .def_static("get_outputs_name", &MaxPooling_Op<DIM>::getOutputsName)
@@ -55,14 +59,17 @@ template <DimIdx_t DIM> void declare_MaxPoolingOp(py::module &m) {
   m.def(("MaxPooling" + std::to_string(DIM) + "D").c_str(), [](const std::vector<DimSize_t>& kernel_dims,
                                                                   const std::string& name,
                                                                   const std::vector<DimSize_t> &stride_dims,
+                                                                  const std::vector<DimSize_t> &dilations,
                                                                   bool ceil_mode) {
         AIDGE_ASSERT(kernel_dims.size() == DIM, "kernel_dims size [{}] does not match DIM [{}]", kernel_dims.size(), DIM);
         AIDGE_ASSERT(stride_dims.size() == DIM, "stride_dims size [{}] does not match DIM [{}]", stride_dims.size(), DIM);
+        AIDGE_ASSERT(dilations.size() == DIM, "dilations size [{}] does not match DIM [{}]", dilations.size(), DIM);
 
-        return MaxPooling<DIM>(to_array<DIM>(kernel_dims.begin()), name, to_array<DIM>(stride_dims.begin()), ceil_mode);
+        return MaxPooling<DIM>(to_array<DIM>(kernel_dims.begin()), name, to_array<DIM>(stride_dims.begin()),  to_array<DIM>(dilations.begin()), ceil_mode);
     }, py::arg("kernel_dims"),
        py::arg("name") = "",
        py::arg("stride_dims") = std::vector<DimSize_t>(DIM, 1),
+       py::arg("dilations") = std::vector<DimSize_t>(DIM, 1),
        py::arg("ceil_mode") = false,
     R"mydelimiter(
         Initialize a node containing a MaxPooling operator.
@@ -75,6 +82,8 @@ template <DimIdx_t DIM> void declare_MaxPoolingOp(py::module &m) {
         :type kernel_dims: List[int]
         :param stride_dims: The stride (step size) to move the kernel over the input.
         :type stride_dims: List[int]
+        :param dilations: The dilation value along each spatial axis of filter.
+        :type dilations: List[int]
         :param ceil_mode: Whether to use ceil or floor when calculating the output dimensions.
         :type ceil_mode: bool
         :param name: Name of the node (optional).
diff --git a/src/operator/MaxPooling.cpp b/src/operator/MaxPooling.cpp
index 535b53749..afd8e00cc 100644
--- a/src/operator/MaxPooling.cpp
+++ b/src/operator/MaxPooling.cpp
@@ -25,11 +25,13 @@ const std::string Aidge::MaxPooling_Op<DIM>::Type = "MaxPooling" + std::to_strin
 template <Aidge::DimIdx_t DIM>
 Aidge::MaxPooling_Op<DIM>::MaxPooling_Op(const std::array<Aidge::DimSize_t, DIM> &kernel_dims,
                             const std::array<Aidge::DimSize_t, DIM> &stride_dims,
+                            const std::array<Aidge::DimSize_t, DIM> &dilations,
                             bool ceil_mode)
     : OperatorTensor(Type, {InputCategory::Data}, 1),
     mAttributes(std::make_shared<Attributes_>(
     attr<MaxPoolingAttr::StrideDims>(stride_dims),
     attr<MaxPoolingAttr::KernelDims>(kernel_dims),
+    attr<MaxPoolingAttr::Dilations>(dilations),
     attr<MaxPoolingAttr::CeilMode>(ceil_mode)))
 {}
 
@@ -63,11 +65,15 @@ bool Aidge::MaxPooling_Op<DIM>::forwardDims(bool /*allowDataDependency*/) {
             roundingFunction = [](float x) { return std::floor(x); };
         }
 
-        for (std::size_t dim = 0; dim < mAttributes->template getAttr<MaxPoolingAttr::KernelDims>().size() ; ++dim) {
+        for (std::size_t dim = 0; dim < mAttributes->template getAttr<MaxPoolingAttr::KernelDims>().size(); ++dim) {
+            const auto kernelDim = mAttributes->template getAttr<MaxPoolingAttr::KernelDims>()[dim];
+            const auto strideDim = mAttributes->template getAttr<MaxPoolingAttr::StrideDims>()[dim];
+            const auto dilationDim = mAttributes->template getAttr<MaxPoolingAttr::Dilations>()[dim];
+
             outputDims[dim+2] = 1 + static_cast<DimSize_t>(
-                                        roundingFunction(static_cast<float>(inputDims[dim+2] -
-                                                                mAttributes->template getAttr<MaxPoolingAttr::KernelDims>()[dim]) /
-                                        static_cast<float>(mAttributes->template getAttr<MaxPoolingAttr::StrideDims>()[dim])));
+                                            roundingFunction(static_cast<float>(inputDims[dim+2] -
+                                                                    (kernelDim - 1) * dilationDim - 1) /
+                                            static_cast<float>(strideDim)));
         }
         outputDims[1] = inputDims[1];
         outputDims[0] = inputDims[0];
@@ -98,12 +104,13 @@ template <std::array<Aidge::DimSize_t, 1>::size_type DIM>
 std::shared_ptr<Aidge::Node> Aidge::MaxPooling(const std::array<Aidge::DimSize_t, DIM> &kernel_dims,
                                            const std::string& name,
                                            const std::array<Aidge::DimSize_t, DIM> &stride_dims,
+                                           const std::array<Aidge::DimSize_t, DIM> &dilations,
                                            bool ceil_mode)
 {
     static_assert(DIM<=MaxDim,"Too many kernel dimensions required by MaxPooling, not supported");
-    return std::make_shared<Node>(std::make_shared<MaxPooling_Op<static_cast<DimIdx_t>(DIM)>>(kernel_dims, stride_dims, ceil_mode), name);
+    return std::make_shared<Node>(std::make_shared<MaxPooling_Op<static_cast<DimIdx_t>(DIM)>>(kernel_dims, stride_dims, dilations, ceil_mode), name);
 }
 
-template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<1>(const std::array<Aidge::DimSize_t, 1>&, const std::string&, const std::array<Aidge::DimSize_t, 1>&, bool);
-template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<2>(const std::array<Aidge::DimSize_t, 2>&, const std::string&, const std::array<Aidge::DimSize_t, 2>&, bool);
-template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<3>(const std::array<Aidge::DimSize_t, 3>&, const std::string&, const std::array<Aidge::DimSize_t, 3>&, bool);
+template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<1>(const std::array<Aidge::DimSize_t, 1>&, const std::string&, const std::array<Aidge::DimSize_t, 1>&, const std::array<Aidge::DimSize_t, 1>&, bool);
+template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<2>(const std::array<Aidge::DimSize_t, 2>&, const std::string&, const std::array<Aidge::DimSize_t, 2>&, const std::array<Aidge::DimSize_t, 2>&, bool);
+template std::shared_ptr<Aidge::Node> Aidge::MaxPooling<3>(const std::array<Aidge::DimSize_t, 3>&, const std::string&, const std::array<Aidge::DimSize_t, 3>&, const std::array<Aidge::DimSize_t, 3>&, bool);
-- 
GitLab