Generic flatten() with 2-d and 3-d inputs

Author: milancurcic

src/nf/nf_flatten_layer.f90

@@ -18,13 +18,20 @@ module nf_flatten_layer
     integer, allocatable :: input_shape(:)
     integer :: output_size
 
-    real, allocatable :: gradient(:,:,:)
+    real, allocatable :: gradient_2d(:,:)
+    real, allocatable :: gradient_3d(:,:,:)
     real, allocatable :: output(:)
 
   contains
 
-    procedure :: backward
-    procedure :: forward
+    procedure :: backward_2d
+    procedure :: backward_3d
+    generic :: backward => backward_2d, backward_3d
+
+    procedure :: forward_2d
+    procedure :: forward_3d
+    generic :: forward => forward_2d, forward_3d
+
     procedure :: init
 
   end type flatten_layer
@@ -39,26 +46,47 @@ end function flatten_layer_cons
 
   interface
 
-    pure module subroutine backward(self, input, gradient)
-      !! Apply the backward pass to the flatten layer.
+    pure module subroutine backward_2d(self, input, gradient)
+      !! Apply the backward pass to the flatten layer for 2D input.
+      !! This is a reshape operation from 1-d gradient to 2-d input.
+      class(flatten_layer), intent(in out) :: self
+        !! Flatten layer instance
+      real, intent(in) :: input(:,:)
+        !! Input from the previous layer
+      real, intent(in) :: gradient(:)
+        !! Gradient from the next layer
+    end subroutine backward_2d
+
+    pure module subroutine backward_3d(self, input, gradient)
+      !! Apply the backward pass to the flatten layer for 3D input.
       !! This is a reshape operation from 1-d gradient to 3-d input.
       class(flatten_layer), intent(in out) :: self
         !! Flatten layer instance
       real, intent(in) :: input(:,:,:)
         !! Input from the previous layer
       real, intent(in) :: gradient(:)
         !! Gradient from the next layer
-    end subroutine backward
+    end subroutine backward_3d
+
+    pure module subroutine forward_2d(self, input)
+      !! Propagate forward the layer for 2D input.
+      !! Calling this subroutine updates the values of a few data components
+      !! of `flatten_layer` that are needed for the backward pass.
+      class(flatten_layer), intent(in out) :: self
+        !! Dense layer instance
+      real, intent(in) :: input(:,:)
+        !! Input from the previous layer
+    end subroutine forward_2d
 
-    pure module subroutine forward(self, input)
-      !! Propagate forward the layer.
+    pure module subroutine forward_3d(self, input)
+      !! Propagate forward the layer for 3D input.
       !! Calling this subroutine updates the values of a few data components
       !! of `flatten_layer` that are needed for the backward pass.
       class(flatten_layer), intent(in out) :: self
         !! Dense layer instance
       real, intent(in) :: input(:,:,:)
         !! Input from the previous layer
-    end subroutine forward
+    end subroutine forward_3d
 
     module subroutine init(self, input_shape)
       !! Initialize the layer data structures.

src/nf/nf_flatten_layer_submodule.f90

@@ -15,19 +15,34 @@ elemental module function flatten_layer_cons() result(res)
   end function flatten_layer_cons
 
 
-  pure module subroutine backward(self, input, gradient)
+  pure module subroutine backward_2d(self, input, gradient)
+    class(flatten_layer), intent(in out) :: self
+    real, intent(in) :: input(:,:)
+    real, intent(in) :: gradient(:)
+    self % gradient_2d = reshape(gradient, shape(input))
+  end subroutine backward_2d
+
+
+  pure module subroutine backward_3d(self, input, gradient)
     class(flatten_layer), intent(in out) :: self
     real, intent(in) :: input(:,:,:)
     real, intent(in) :: gradient(:)
-    self % gradient = reshape(gradient, shape(input))
-  end subroutine backward
+    self % gradient_3d = reshape(gradient, shape(input))
+  end subroutine backward_3d
+
+
+  pure module subroutine forward_2d(self, input)
+    class(flatten_layer), intent(in out) :: self
+    real, intent(in) :: input(:,:)
+    self % output = pack(input, .true.)
+  end subroutine forward_2d
 
 
-  pure module subroutine forward(self, input)
+  pure module subroutine forward_3d(self, input)
     class(flatten_layer), intent(in out) :: self
     real, intent(in) :: input(:,:,:)
     self % output = pack(input, .true.)
-  end subroutine forward
+  end subroutine forward_3d
 
 
   module subroutine init(self, input_shape)
@@ -37,8 +52,13 @@ module subroutine init(self, input_shape)
     self % input_shape = input_shape
     self % output_size = product(input_shape)
 
-    allocate(self % gradient(input_shape(1), input_shape(2), input_shape(3)))
-    self % gradient = 0
+    if (size(input_shape) == 2) then
+      allocate(self % gradient_2d(input_shape(1), input_shape(2)))
+      self % gradient_2d = 0
+    else if (size(input_shape) == 3) then
+      allocate(self % gradient_3d(input_shape(1), input_shape(2), input_shape(3)))
+      self % gradient_3d = 0
+    end if
 
     allocate(self % output(self % output_size))
     self % output = 0

src/nf/nf_layer_submodule.f90

@@ -37,8 +37,10 @@ pure module subroutine backward_1d(self, previous, gradient)
 
       type is(flatten_layer)
 
-        ! Upstream layers permitted: input3d, conv2d, maxpool2d
+        ! Upstream layers permitted: input2d, input3d, conv2d, maxpool2d
         select type(prev_layer => previous % p)
+          type is(input2d_layer)
+            call this_layer % backward(prev_layer % output, gradient)
           type is(input3d_layer)
             call this_layer % backward(prev_layer % output, gradient)
           type is(conv2d_layer)
@@ -168,8 +170,10 @@ pure module subroutine forward(self, input)
 
       type is(flatten_layer)
 
-        ! Upstream layers permitted: input3d, conv2d, maxpool2d, reshape3d
+        ! Upstream layers permitted: input2d, input3d, conv2d, maxpool2d, reshape3d
         select type(prev_layer => input % p)
+          type is(input2d_layer)
+            call this_layer % forward(prev_layer % output)
           type is(input3d_layer)
             call this_layer % forward(prev_layer % output)
           type is(conv2d_layer)

src/nf/nf_network_submodule.f90

@@ -135,12 +135,20 @@ module subroutine backward(self, output, loss)
         select type(next_layer => self % layers(n + 1) % p)
           type is(dense_layer)
             call self % layers(n) % backward(self % layers(n - 1), next_layer % gradient)
+
           type is(conv2d_layer)
             call self % layers(n) % backward(self % layers(n - 1), next_layer % gradient)
+
           type is(flatten_layer)
-            call self % layers(n) % backward(self % layers(n - 1), next_layer % gradient)
+            if (size(self % layers(n) % layer_shape) == 2) then
+              call self % layers(n) % backward(self % layers(n - 1), next_layer % gradient_2d)
+            else
+              call self % layers(n) % backward(self % layers(n - 1), next_layer % gradient_3d)
+            end if
+
           type is(maxpool2d_layer)
             call self % layers(n) % backward(self % layers(n - 1), next_layer % gradient)
+
           type is(reshape3d_layer)
             call self % layers(n) % backward(self % layers(n - 1), next_layer % gradient)
         end select

test/test_flatten_layer.f90

@@ -3,16 +3,18 @@ program test_flatten_layer
   use iso_fortran_env, only: stderr => error_unit
   use nf, only: dense, flatten, input, layer, network
   use nf_flatten_layer, only: flatten_layer
+  use nf_input2d_layer, only: input2d_layer
   use nf_input3d_layer, only: input3d_layer
 
   implicit none
 
   type(layer) :: test_layer, input_layer
   type(network) :: net
-  real, allocatable :: gradient(:,:,:)
+  real, allocatable :: gradient_3d(:,:,:), gradient_2d(:,:)
   real, allocatable :: output(:)
   logical :: ok = .true.
 
+  ! Test 3D input
   test_layer = flatten()
 
   if (.not. test_layer % name == 'flatten') then
@@ -59,14 +61,49 @@ program test_flatten_layer
   call test_layer % backward(input_layer, real([1, 2, 3, 4]))
 
   select type(this_layer => test_layer % p); type is(flatten_layer)
-    gradient = this_layer % gradient
+    gradient_3d = this_layer % gradient_3d
   end select
 
-  if (.not. all(gradient == reshape(real([1, 2, 3, 4]), [1, 2, 2]))) then
+  if (.not. all(gradient_3d == reshape(real([1, 2, 3, 4]), [1, 2, 2]))) then
     ok = .false.
     write(stderr, '(a)') 'flatten layer correctly propagates backward.. failed'
   end if
 
+  ! Test 2D input
+  test_layer = flatten()
+  input_layer = input(2, 3)
+  call test_layer % init(input_layer)
+
+  if (.not. all(test_layer % layer_shape == [6])) then
+    ok = .false.
+    write(stderr, '(a)') 'flatten layer has an incorrect output shape for 2D input.. failed'
+  end if
+
+  ! Test forward pass - reshaping from 2-d to 1-d
+  select type(this_layer => input_layer % p); type is(input2d_layer)
+    call this_layer % set(reshape(real([1, 2, 3, 4, 5, 6]), [2, 3]))
+  end select
+
+  call test_layer % forward(input_layer)
+  call test_layer % get_output(output)
+
+  if (.not. all(output == [1, 2, 3, 4, 5, 6])) then
+    ok = .false.
+    write(stderr, '(a)') 'flatten layer correctly propagates forward for 2D input.. failed'
+  end if
+
+  ! Test backward pass - reshaping from 1-d to 2-d
+  call test_layer % backward(input_layer, real([1, 2, 3, 4, 5, 6]))
+
+  select type(this_layer => test_layer % p); type is(flatten_layer)
+    gradient_2d = this_layer % gradient_2d
+  end select
+
+  if (.not. all(gradient_2d == reshape(real([1, 2, 3, 4, 5, 6]), [2, 3]))) then
+    ok = .false.
+    write(stderr, '(a)') 'flatten layer correctly propagates backward for 2D input.. failed'
+  end if
+
   net = network([ &
     input(1, 28, 28), &
     flatten(), &