Poolings

CMakeLists.txt

@@ -17,6 +17,12 @@ include(cmake/compilers.cmake)
 add_library(neural-fortran
   src/nf.f90
   src/nf/nf_activation.f90
+  src/nf/nf_avgpool1d_layer.f90
+  src/nf/nf_avgpool1d_layer_submodule.f90
+  src/nf/nf_avgpool2d_layer.f90
+  src/nf/nf_avgpool2d_layer_submodule.f90
+  src/nf/nf_avgpool3d_layer.f90
+  src/nf/nf_avgpool3d_layer_submodule.f90
   src/nf/nf_base_layer.f90
   src/nf/nf_conv1d_layer.f90
   src/nf/nf_conv1d_layer_submodule.f90
@@ -55,6 +61,8 @@ add_library(neural-fortran
   src/nf/nf_maxpool1d_layer_submodule.f90
   src/nf/nf_maxpool2d_layer.f90
   src/nf/nf_maxpool2d_layer_submodule.f90
+  src/nf/nf_maxpool3d_layer.f90
+  src/nf/nf_maxpool3d_layer_submodule.f90
   src/nf/nf_metrics.f90
   src/nf/nf_multihead_attention.f90
   src/nf/nf_multihead_attention_submodule.f90

src/nf.f90

@@ -3,6 +3,9 @@ module nf
   use nf_datasets_mnist, only: label_digits, load_mnist
   use nf_layer, only: layer
   use nf_layer_constructors, only: &
+    avgpool1d, &
+    avgpool2d, &
+    avgpool3d, &
     conv1d, &
     conv2d, &
     dense, &
@@ -15,6 +18,7 @@ module nf
     locally_connected1d, &
     maxpool1d, &
     maxpool2d, &
+    maxpool3d, &
     reshape, &
     self_attention
   use nf_loss, only: mse, quadratic

src/nf/nf_avgpool1d_layer.f90

@@ -0,0 +1,66 @@
+module nf_avgpool1d_layer
+    !! This module provides the 1-d average pooling layer.
+
+    use nf_base_layer, only: base_layer
+    implicit none
+
+    private
+    public :: avgpool1d_layer
+
+    type, extends(base_layer) :: avgpool1d_layer
+        integer :: channels
+        integer :: width      ! Length of the input along the pooling dimension
+        integer :: pool_size
+        integer :: stride
+
+        ! Gradient for the input (same shape as the input).
+        real, allocatable :: gradient(:,:)
+        ! Output after pooling (dimensions: (channels, new_width)).
+        real, allocatable :: output(:,:)
+    contains
+        procedure :: init
+        procedure :: forward
+        procedure :: backward
+    end type avgpool1d_layer
+
+    interface avgpool1d_layer
+        pure module function avgpool1d_layer_cons(pool_size, stride) result(res)
+            !! `avgpool1d` constructor function.
+            integer, intent(in) :: pool_size
+                !! Width of the pooling window.
+            integer, intent(in) :: stride
+                !! Stride of the pooling window.
+            type(avgpool1d_layer) :: res
+        end function avgpool1d_layer_cons
+    end interface avgpool1d_layer
+
+    interface
+        module subroutine init(self, input_shape)
+            !! Initialize the `avgpool1d` layer instance with an input shape.
+            class(avgpool1d_layer), intent(in out) :: self
+                !! `avgpool1d_layer` instance.
+            integer, intent(in) :: input_shape(:)
+                !! Array shape of the input layer, expected as (channels, width).
+        end subroutine init
+
+        pure module subroutine forward(self, input)
+            !! Run a forward pass of the `avgpool1d` layer.
+            class(avgpool1d_layer), intent(in out) :: self
+                !! `avgpool1d_layer` instance.
+            real, intent(in) :: input(:,:)
+                !! Input data (output of the previous layer), with shape (channels, width).
+        end subroutine forward
+
+        pure module subroutine backward(self, input, gradient)
+            !! Run a backward pass of the `avgpool1d` layer.
+            class(avgpool1d_layer), intent(in out) :: self
+                !! `avgpool1d_layer` instance.
+            real, intent(in) :: input(:,:)
+                !! Input data (output of the previous layer).
+            real, intent(in) :: gradient(:,:)
+                !! Gradient from the downstream layer, with shape (channels, pooled width).
+        end subroutine backward
+    end interface
+
+  end module nf_avgpool1d_layer
+

src/nf/nf_avgpool1d_layer_submodule.f90

@@ -0,0 +1,87 @@
+submodule(nf_avgpool1d_layer) nf_avgpool1d_layer_submodule
+  implicit none
+
+contains
+
+  pure module function avgpool1d_layer_cons(pool_size, stride) result(res)
+    implicit none
+    integer, intent(in) :: pool_size
+    integer, intent(in) :: stride
+    type(avgpool1d_layer) :: res
+
+    res % pool_size = pool_size
+    res % stride    = stride
+  end function avgpool1d_layer_cons
+
+
+  module subroutine init(self, input_shape)
+    implicit none
+    class(avgpool1d_layer), intent(in out) :: self
+    integer, intent(in) :: input_shape(:)
+    ! input_shape is expected to be (channels, width)
+
+    self % channels = input_shape(1)
+    ! The new width is the integer division of the input width by the stride.
+    self % width    = input_shape(2) / self % stride
+
+    ! Allocate the gradient array corresponding to the input dimensions.
+    allocate(self % gradient(input_shape(1), input_shape(2)))
+    self % gradient = 0
+
+    ! Allocate the output array (after pooling).
+    allocate(self % output(self % channels, self % width))
+    self % output = 0
+  end subroutine init
+
+
+  pure module subroutine forward(self, input)
+    implicit none
+    class(avgpool1d_layer), intent(in out) :: self
+    real, intent(in) :: input(:,:)
+    integer :: input_width
+    integer :: i, n
+    integer :: ii, iend
+    integer :: iextent
+
+    input_width = size(input, dim=2)
+    ! Ensure we only process complete pooling regions.
+    iextent = input_width - mod(input_width, self % stride)
+
+    ! Loop over the input with a step size equal to the stride and over all channels.
+    do concurrent (i = 1:iextent: self % stride, n = 1:self % channels)
+      ! Compute the index in the pooled (output) array.
+      ii = (i - 1) / self % stride + 1
+      ! Determine the ending index of the current pooling region.
+      iend = min(i + self % pool_size - 1, input_width)
+
+      ! Compute the average over the pooling region.
+      self % output(n, ii) = sum(input(n, i:iend)) / (iend - i + 1)
+    end do
+  end subroutine forward
+
+
+  pure module subroutine backward(self, input, gradient)
+    implicit none
+    class(avgpool1d_layer), intent(in out) :: self
+    real, intent(in) :: input(:,:)
+    real, intent(in) :: gradient(:,:)
+    integer :: channels, pooled_width
+    integer :: i, n, j, istart, iend
+    real :: scale_factor
+
+    channels    = size(gradient, dim=1)
+    pooled_width = size(gradient, dim=2)
+
+    ! The gradient for average pooling is distributed evenly over the pooling window.
+    do concurrent (n = 1:channels, i = 1:pooled_width)
+      istart = (i - 1) * self % stride + 1
+      iend = min(istart + self % pool_size - 1, size(input, dim=2))
+      scale_factor = 1.0 / (iend - istart + 1)
+
+      do j = istart, iend
+        self % gradient(n, j) = gradient(n, i) * scale_factor
+      end do
+    end do
+  end subroutine backward
+
+end submodule nf_avgpool1d_layer_submodule

src/nf/nf_avgpool2d_layer.f90

@@ -0,0 +1,66 @@
+module nf_avgpool2d_layer
+    !! This module provides the 2-d average pooling layer.
+
+    use nf_base_layer, only: base_layer
+    implicit none
+
+    private
+    public :: avgpool2d_layer
+
+    type, extends(base_layer) :: avgpool2d_layer
+        integer :: channels
+        integer :: height      ! Height of the input
+        integer :: width       ! Width of the input
+        integer :: pool_size  ! Pooling window size (height, width)
+        integer :: stride    ! Stride (height, width)
+
+        ! Gradient for the input (same shape as the input: channels, height, width).
+        real, allocatable :: gradient(:,:,:)
+        ! Output after pooling (dimensions: (channels, new_height, new_width)).
+        real, allocatable :: output(:,:,:)
+    contains
+        procedure :: init
+        procedure :: forward
+        procedure :: backward
+    end type avgpool2d_layer
+
+    interface avgpool2d_layer
+        pure module function avgpool2d_layer_cons(pool_size, stride) result(res)
+            !! `avgpool2d` constructor function.
+            integer, intent(in) :: pool_size
+                !! Pooling window size (height, width).
+            integer, intent(in) :: stride
+                !! Stride (height, width).
+            type(avgpool2d_layer) :: res
+        end function avgpool2d_layer_cons
+    end interface avgpool2d_layer
+
+    interface
+        module subroutine init(self, input_shape)
+            !! Initialize the `avgpool2d` layer instance with an input shape.
+            class(avgpool2d_layer), intent(in out) :: self
+                !! `avgpool2d_layer` instance.
+            integer, intent(in) :: input_shape(:)
+                !! Array shape of the input layer, expected as (channels, height, width).
+        end subroutine init
+
+        pure module subroutine forward(self, input)
+            !! Run a forward pass of the `avgpool2d` layer.
+            class(avgpool2d_layer), intent(in out) :: self
+                !! `avgpool2d_layer` instance.
+            real, intent(in) :: input(:,:,:)
+                !! Input data (output of the previous layer), with shape (channels, height, width).
+        end subroutine forward
+
+        pure module subroutine backward(self, input, gradient)
+            !! Run a backward pass of the `avgpool2d` layer.
+            class(avgpool2d_layer), intent(in out) :: self
+                !! `avgpool2d_layer` instance.
+            real, intent(in) :: input(:,:,:)
+                !! Input data (output of the previous layer).
+            real, intent(in) :: gradient(:,:,:)
+                !! Gradient from the downstream layer, with shape (channels, pooled_height, pooled_width).
+        end subroutine backward
+    end interface
+
+end module nf_avgpool2d_layer

src/nf/nf_avgpool2d_layer_submodule.f90

@@ -0,0 +1,94 @@
+submodule(nf_avgpool2d_layer) nf_avgpool2d_layer_submodule
+  implicit none
+
+contains
+
+  pure module function avgpool2d_layer_cons(pool_size, stride) result(res)
+    implicit none
+    integer, intent(in) :: pool_size
+    integer, intent(in) :: stride
+    type(avgpool2d_layer) :: res
+
+    res % pool_size = pool_size
+    res % stride    = stride
+  end function avgpool2d_layer_cons
+
+
+  module subroutine init(self, input_shape)
+    implicit none
+    class(avgpool2d_layer), intent(in out) :: self
+    integer, intent(in) :: input_shape(:)
+    ! input_shape is expected to be (channels, width, height)
+
+    self % channels = input_shape(1)
+    self % width    = input_shape(2) / self % stride
+    self % height   = input_shape(3) / self % stride
+
+    ! Allocate the gradient array corresponding to the input dimensions.
+    allocate(self % gradient(input_shape(1), input_shape(2), input_shape(3)))
+    self % gradient = 0
+
+    ! Allocate the output array (after pooling).
+    allocate(self % output(self % channels, self % width, self % height))
+    self % output = 0
+  end subroutine init
+
+
+  pure module subroutine forward(self, input)
+    implicit none
+    class(avgpool2d_layer), intent(in out) :: self
+    real, intent(in) :: input(:,:,:)
+    integer :: input_width, input_height
+    integer :: i, j, n
+    integer :: ii, jj, iend, jend
+    integer :: iextent, jextent
+
+    input_width  = size(input, dim=2)
+    input_height = size(input, dim=3)
+
+    ! Ensure we only process complete pooling regions.
+    iextent = input_width - mod(input_width, self % stride)
+    jextent = input_height - mod(input_height, self % stride)
+
+    ! Loop over the input with a step size equal to the stride and over all channels.
+    do concurrent (i = 1:iextent:self % stride, j = 1:jextent:self % stride, n = 1:self % channels)
+      ii = (i - 1) / self % stride + 1
+      jj = (j - 1) / self % stride + 1
+
+      iend = min(i + self % pool_size - 1, input_width)
+      jend = min(j + self % pool_size - 1, input_height)
+
+      ! Compute the average over the pooling region.
+      self % output(n, ii, jj) = sum(input(n, i:iend, j:jend)) / ((iend - i + 1) * (jend - j + 1))
+    end do
+  end subroutine forward
+
+
+  pure module subroutine backward(self, input, gradient)
+    implicit none
+    class(avgpool2d_layer), intent(in out) :: self
+    real, intent(in) :: input(:,:,:)
+    real, intent(in) :: gradient(:,:,:)
+    integer :: channels, pooled_width, pooled_height
+    integer :: i, j, n, x, y, istart, iend, jstart, jend
+    real :: scale_factor
+
+    channels      = size(gradient, dim=1)
+    pooled_width  = size(gradient, dim=2)
+    pooled_height = size(gradient, dim=3)
+
+    ! The gradient for average pooling is distributed evenly over the pooling window.
+    do concurrent (n = 1:channels, i = 1:pooled_width, j = 1:pooled_height)
+      istart = (i - 1) * self % stride + 1
+      iend   = min(istart + self % pool_size - 1, size(input, dim=2))
+      jstart = (j - 1) * self % stride + 1
+      jend   = min(jstart + self % pool_size - 1, size(input, dim=3))
+      scale_factor = 1.0 / ((iend - istart + 1) * (jend - jstart + 1))
+
+      do concurrent (x = istart:iend, y = jstart:jend)
+        self % gradient(n, x, y) = gradient(n, i, j) * scale_factor
+      end do
+    end do
+  end subroutine backward
+
+end submodule nf_avgpool2d_layer_submodule