Forward pass for a max-pooling layer

CMakeLists.txt

@@ -86,6 +86,8 @@ add_library(neural
   src/nf_layer_submodule.f90
   src/nf_loss.f90
   src/nf_loss_submodule.f90
+  src/nf_maxpool2d_layer.f90
+  src/nf_maxpool2d_layer_submodule.f90
   src/nf_network.f90
   src/nf_network_submodule.f90
   src/nf_optimizers.f90
@@ -100,7 +102,7 @@ string(REGEX REPLACE "^ | $" "" LIBS "${LIBS}")
 
 # tests
 enable_testing()
-foreach(execid input1d_layer input3d_layer dense_layer conv2d_layer dense_network conv2d_network)
+foreach(execid input1d_layer input3d_layer dense_layer conv2d_layer maxpool2d_layer dense_network conv2d_network)
   add_executable(test_${execid} test/test_${execid}.f90)
   target_link_libraries(test_${execid} neural ${LIBS})
   add_test(test_${execid} bin/test_${execid})

README.md

@@ -16,18 +16,20 @@ Read the paper [here](https://arxiv.org/abs/1902.06714).
 
 ## Features
 
-* Dense, fully connected neural networks of arbitrary shape and size
-* Backprop with Mean Square Error cost function
+* Dense, fully connected neural layers
+* Convolutional and max-pooling layers (experimental, forward propagation only)
+* Stochastic and mini-batch gradient descent for back-propagation
 * Data-based parallelism
 * Several activation functions
 
 ### Available layer types
 
-| Layer type | Constructor name | Rank of output array | Forward pass | Backward pass |
-|------------|------------------|----------------------|--------------|---------------|
-| Input | `input` | 1, 3 | n/a | n/a |
-| Dense (fully-connected) | `dense` | 1 | ✅ | ✅ |
-| Convolutional (2-d) | `conv2d` | 3 | ✅ | ❌ |
+| Layer type | Constructor name | Supported input layers | Rank of output array | Forward pass | Backward pass |
+|------------|------------------|------------------------|----------------------|--------------|---------------|
+| Input (1-d and 3-d) | `input` | n/a | 1, 3 | n/a | n/a |
+| Dense (fully-connected) | `input` (1-d) | `dense` | 1 | ✅ | ✅ |
+| Convolutional (2-d) | `input` (3-d), `conv2d`, `maxpool2d` | `conv2d` | 3 | ✅ | ❌ |
+| Max-pooling (2-d) | `input` (3-d), `conv2d`, `maxpool2d` | `maxpool2d` | 3 | ✅ | ❌ |
 
 ## Getting started
 

src/nf.f90

@@ -1,6 +1,6 @@
 module nf
   use nf_datasets_mnist, only: label_digits, load_mnist
   use nf_layer, only: layer
-  use nf_layer_constructors, only: conv2d, dense, input
+  use nf_layer_constructors, only: conv2d, dense, input, maxpool2d
   use nf_network, only: network
 end module nf

src/nf_layer_constructors.f90

@@ -7,7 +7,7 @@ module nf_layer_constructors
   implicit none
 
   private
-  public :: conv2d, dense, input
+  public :: conv2d, dense, input, maxpool2d
 
   interface input
 
@@ -111,6 +111,28 @@ pure module function conv2d(filters, kernel_size, activation) result(res)
         !! Resulting layer instance
     end function conv2d
 
+    pure module function maxpool2d(pool_size, stride) result(res)
+      !! 2-d maxpooling layer constructor.
+      !!
+      !! This layer is for downscaling other layers, typically `conv2d`.
+      !!
+      !! Example:
+      !!
+      !! ```
+      !! use nf, only :: maxpool2d, layer
+      !! type(layer) :: maxpool2d_layer
+      !! maxpool2d_layer = maxpool2d(pool_size=2)
+      !! maxpool2d_layer = maxpool2d(pool_size=2, stride=3)
+      !! ```
+      integer, intent(in) :: pool_size
+        !! Width of the pooling window, commonly 2
+      integer, intent(in), optional :: stride
+        !! Stride of the pooling window, commonly equal to `pool_size`;
+        !! Defaults to `pool_size` if omitted.
+      type(layer) :: res
+        !! Resulting layer instance
+    end function maxpool2d
+
   end interface
 
 end module nf_layer_constructors

src/nf_layer_constructors_submodule.f90

@@ -5,6 +5,7 @@
   use nf_dense_layer, only: dense_layer
   use nf_input1d_layer, only: input1d_layer
   use nf_input3d_layer, only: input3d_layer
+  use nf_maxpool2d_layer, only: maxpool2d_layer
 
   implicit none
 
@@ -72,4 +73,33 @@ pure module function conv2d(filters, kernel_size, activation) result(res)
 
   end function conv2d
 
+
+  pure module function maxpool2d(pool_size, stride) result(res)
+    integer, intent(in) :: pool_size
+    integer, intent(in), optional :: stride
+    integer :: stride_
+    type(layer) :: res
+
+    if (pool_size < 2) &
+      error stop 'pool_size must be >= 2 in a maxpool2d layer'
+
+    ! Stride defaults to pool_size if not provided
+    if (present(stride)) then
+      stride_ = stride
+    else
+      stride_ = pool_size
+    end if
+
+    if (stride_ < 1) &
+      error stop 'stride must be >= 1 in a maxpool2d layer'
+
+    res % name = 'maxpool2d'
+
+    allocate( &
+      res % p, &
+      source=maxpool2d_layer(pool_size, stride_) &
+    )
+
+  end function maxpool2d
+
 end submodule nf_layer_constructors_submodule

src/nf_layer_submodule.f90

@@ -4,6 +4,7 @@
   use nf_dense_layer, only: dense_layer
   use nf_input1d_layer, only: input1d_layer
   use nf_input3d_layer, only: input3d_layer
+  use nf_maxpool2d_layer, only: maxpool2d_layer
 
   implicit none
 
@@ -51,12 +52,26 @@ pure module subroutine forward(self, input)
 
       type is(conv2d_layer)
 
-        ! Input layers permitted: input3d, conv2d
+        ! Input layers permitted: input3d, conv2d, maxpool2d
         select type(prev_layer => input % p)
           type is(input3d_layer)
             call this_layer % forward(prev_layer % output)
           type is(conv2d_layer)
             call this_layer % forward(prev_layer % output)
+          type is(maxpool2d_layer)
+            call this_layer % forward(prev_layer % output)
+        end select
+
+      type is(maxpool2d_layer)
+
+        ! Input layers permitted: input3d, conv2d, maxpool2d
+        select type(prev_layer => input % p)
+          type is(input3d_layer)
+            call this_layer % forward(prev_layer % output)
+          type is(conv2d_layer)
+            call this_layer % forward(prev_layer % output)
+          type is(maxpool2d_layer)
+            call this_layer % forward(prev_layer % output)
         end select
 
     end select
@@ -92,8 +107,10 @@ pure module subroutine get_output_3d(self, output)
         allocate(output, source=this_layer % output)
       type is(conv2d_layer)
         allocate(output, source=this_layer % output)
+      type is(maxpool2d_layer)
+        allocate(output, source=this_layer % output)
       class default
-        error stop '3-d output can only be read from an input3d or conv2d layer.'
+        error stop '3-d output can only be read from an input3d, conv2d, or maxpool2d layer.'
 
     end select
 
@@ -111,9 +128,13 @@ impure elemental module subroutine init(self, input)
       call this_layer % init(input % layer_shape)
     end select
 
-    ! The shape of a conv2d layer is not known until we receive an input layer.
-    select type(this_layer => self % p); type is(conv2d_layer)
-      self % layer_shape = shape(this_layer % output)
+    ! The shape of conv2d or maxpool2d layers is not known
+    ! until we receive an input layer.
+    select type(this_layer => self % p)
+      type is(conv2d_layer)
+        self % layer_shape = shape(this_layer % output)
+      type is(maxpool2d_layer)
+        self % layer_shape = shape(this_layer % output)
     end select
 
     self % input_layer_shape = input % layer_shape 

src/nf_maxpool2d_layer.f90

@@ -0,0 +1,75 @@
+module nf_maxpool2d_layer
+
+  !! This module provides the 2-d maxpooling layer.
+
+  use nf_base_layer, only: base_layer
+  implicit none
+
+  private
+  public :: maxpool2d_layer
+
+  type, extends(base_layer) :: maxpool2d_layer
+
+    integer :: channels
+    integer :: width
+    integer :: height
+    integer :: pool_size
+    integer :: stride
+
+    ! Locations (as input matrix indices) of the maximum values
+    ! in the width (x) and height (y) dimensions
+    integer, allocatable :: maxloc_x(:,:,:)
+    integer, allocatable :: maxloc_y(:,:,:)
+
+    real, allocatable :: output(:,:,:)
+
+  contains
+
+    procedure :: init
+    procedure :: forward
+    procedure :: backward
+
+  end type maxpool2d_layer
+
+  interface maxpool2d_layer
+    pure module function maxpool2d_layer_cons(pool_size, stride) result(res)
+      !! `maxpool2d` constructor function
+      integer, intent(in) :: pool_size
+        !! Width and height of the pooling window
+      integer, intent(in) :: stride
+        !! Stride of the pooling window
+      type(maxpool2d_layer) :: res
+    end function maxpool2d_layer_cons
+  end interface maxpool2d_layer
+
+  interface
+
+    module subroutine init(self, input_shape)
+      !! Initialize the `maxpool2d` layer instance with an input shape.
+      class(maxpool2d_layer), intent(in out) :: self
+        !! `maxpool2d_layer` instance
+      integer, intent(in) :: input_shape(:)
+        !! Array shape of the input layer
+    end subroutine init
+
+    pure module subroutine forward(self, input)
+      !! Run a forward pass of the `maxpool2d` layer.
+      class(maxpool2d_layer), intent(in out) :: self
+        !! `maxpool2d_layer` instance
+      real, intent(in) :: input(:,:,:)
+        !! Input data (output of the previous layer)
+    end subroutine forward
+
+    module subroutine backward(self, input, gradient)
+      !! Run a backward pass of the `maxpool2d` layer.
+      class(maxpool2d_layer), intent(in out) :: self
+        !! `maxpool2d_layer` instance
+      real, intent(in) :: input(:,:,:)
+        !! Input data (output of the previous layer)
+      real, intent(in) :: gradient(:,:,:)
+        !! Gradient from the downstream layer
+    end subroutine backward
+
+  end interface
+
+end module nf_maxpool2d_layer

src/nf_maxpool2d_layer_submodule.f90

@@ -0,0 +1,89 @@
+submodule(nf_maxpool2d_layer) nf_maxpool2d_layer_submodule
+
+  implicit none
+
+contains
+
+  pure module function maxpool2d_layer_cons(pool_size, stride) result(res)
+    implicit none
+    integer, intent(in) :: pool_size
+    integer, intent(in) :: stride
+    type(maxpool2d_layer) :: res
+    res % pool_size = pool_size
+    res % stride = stride
+  end function maxpool2d_layer_cons
+
+
+  module subroutine init(self, input_shape)
+    implicit none
+    class(maxpool2d_layer), intent(in out) :: self
+    integer, intent(in) :: input_shape(:)
+
+    self % channels = input_shape(1)
+    self % width = input_shape(2) / self % stride
+    self % height = input_shape(3) / self % stride
+
+    allocate(self % maxloc_x(self % channels, self % width, self % height))
+    self % maxloc_x = 0
+
+    allocate(self % maxloc_y(self % channels, self % width, self % height))
+    self % maxloc_y = 0
+
+    allocate(self % output(self % channels, self % width, self % height))
+    self % output = 0
+
+  end subroutine init
+
+
+  pure module subroutine forward(self, input)
+    implicit none
+    class(maxpool2d_layer), intent(in out) :: self
+    real, intent(in) :: input(:,:,:)
+    integer :: input_width, input_height
+    integer :: i, j, n
+    integer :: ii, jj
+    integer :: iend, jend
+    integer :: maxloc_xy(2)
+
+    input_width = size(input, dim=2)
+    input_height = size(input, dim=2)
+
+    ! Stride along the width and height of the input image
+    stride_over_input: do concurrent( &
+      i = 1:input_width:self % stride, &
+      j = 1:input_height:self % stride &
+    )
+
+      ! Indices of the pooling layer
+      ii = i / self % stride + 1
+      jj = j / self % stride + 1
+
+      iend = i + self % pool_size - 1
+      jend = j + self % pool_size - 1
+
+      maxpool_for_each_channel: do concurrent(n = 1:self % channels)
+
+        ! Get and store the location of the maximum value
+        maxloc_xy = maxloc(input(n,i:iend,j:jend))
+        self % maxloc_x(n,ii,jj) = maxloc_xy(1) + i - 1
+        self % maxloc_y(n,ii,jj) = maxloc_xy(2) + j - 1
+
+        self % output(n,ii,jj) = &
+          input(n,self % maxloc_x(n,ii,jj),self % maxloc_y(n,ii,jj))
+
+      end do maxpool_for_each_channel
+
+    end do stride_over_input
+
+  end subroutine forward
+
+
+  module subroutine backward(self, input, gradient)
+    implicit none
+    class(maxpool2d_layer), intent(in out) :: self
+    real, intent(in) :: input(:,:,:)
+    real, intent(in) :: gradient(:,:,:)
+    print *, 'Warning: maxpool2d backward pass not implemented'
+  end subroutine backward
+
+end submodule nf_maxpool2d_layer_submodule