linear2d_layer: make linear2d layer work with input2d and flatten2d

Author: OneAdder

src/nf.f90

@@ -3,7 +3,7 @@ module nf
   use nf_datasets_mnist, only: label_digits, load_mnist
   use nf_layer, only: layer
   use nf_layer_constructors, only: &
-    conv2d, dense, flatten, input, maxpool2d, reshape, linear2d
+    conv2d, dense, flatten, flatten2d, input, maxpool2d, reshape, linear2d
   use nf_loss, only: mse, quadratic
   use nf_metrics, only: corr, maxabs
   use nf_network, only: network

src/nf/nf_layer_submodule.f90

@@ -4,6 +4,7 @@
   use nf_conv2d_layer, only: conv2d_layer
   use nf_dense_layer, only: dense_layer
   use nf_flatten_layer, only: flatten_layer
+  use nf_flatten2d_layer, only: flatten2d_layer
   use nf_input1d_layer, only: input1d_layer
   use nf_input2d_layer, only: input2d_layer
   use nf_input3d_layer, only: input3d_layer
@@ -48,8 +49,16 @@ pure module subroutine backward_1d(self, previous, gradient)
             call this_layer % backward(prev_layer % output, gradient)
           type is(maxpool2d_layer)
             call this_layer % backward(prev_layer % output, gradient)
-!          type is(linear2d_layer)
-!            call this_layer % backward(prev_layer % output, gradient)
+        end select
+
+      type is(flatten2d_layer)
+
+        ! Upstream layers permitted: linear2d_layer
+        select type(prev_layer => previous % p)
+          type is(linear2d_layer)
+            call this_layer % backward(prev_layer % output, gradient)
+          type is(input2d_layer)
+            call this_layer % backward(prev_layer % output, gradient)
         end select
 
     end select
@@ -63,8 +72,6 @@ pure module subroutine backward_2d(self, previous, gradient)
     class(layer), intent(in) :: previous
     real, intent(in) :: gradient(:,:)
 
-    ! Backward pass from a 2-d layer downstream currently implemented
-    ! only for input2d and linear2d layers
     select type(this_layer => self % p)
 
       type is(linear2d_layer)
@@ -197,8 +204,14 @@ pure module subroutine forward(self, input)
             call this_layer % forward(prev_layer % output)
           type is(reshape3d_layer)
             call this_layer % forward(prev_layer % output)
-!          type is(linear2d_layer)
-!            call this_layer % forward(prev_layer % output)
+        end select
+
+      type is(flatten2d_layer)
+        select type(prev_layer => input % p)
+          type is(linear2d_layer)
+            call this_layer % forward(prev_layer % output)
+          type is(input2d_layer)
+            call this_layer % forward(prev_layer % output)
         end select
 
       type is(reshape3d_layer)
@@ -241,6 +254,8 @@ pure module subroutine get_output_1d(self, output)
         allocate(output, source=this_layer % output)
       type is(flatten_layer)
         allocate(output, source=this_layer % output)
+      type is(flatten2d_layer)
+        allocate(output, source=this_layer % output)
       class default
         error stop '1-d output can only be read from an input1d, dense, or flatten layer.'
 
@@ -312,9 +327,11 @@ impure elemental module subroutine init(self, input)
         self % layer_shape = shape(this_layer % output)
       type is(flatten_layer)
         self % layer_shape = shape(this_layer % output)
+      type is(flatten2d_layer)
+        self % layer_shape = shape(this_layer % output)
     end select
 
-    self % input_layer_shape = input % layer_shape 
+    self % input_layer_shape = input % layer_shape
     self % initialized = .true.
 
   end subroutine init
@@ -355,6 +372,8 @@ elemental module function get_num_params(self) result(num_params)
         num_params = 0
       type is (flatten_layer)
         num_params = 0
+      type is (flatten2d_layer)
+        num_params = 0
       type is (reshape3d_layer)
         num_params = 0
       type is (linear2d_layer)
@@ -384,6 +403,8 @@ module function get_params(self) result(params)
         ! No parameters to get.
       type is (flatten_layer)
         ! No parameters to get.
+      type is (flatten2d_layer)
+        ! No parameters to get.
       type is (reshape3d_layer)
         ! No parameters to get.
       type is (linear2d_layer)
@@ -412,6 +433,8 @@ module function get_gradients(self) result(gradients)
       type is (maxpool2d_layer)
         ! No gradients to get.
       type is (flatten_layer)
+        ! No parameters to get.
+      type is (flatten2d_layer)
         ! No gradients to get.
       type is (reshape3d_layer)
         ! No gradients to get.
@@ -477,6 +500,11 @@ module subroutine set_params(self, params)
         write(stderr, '(a)') 'Warning: calling set_params() ' &
           // 'on a zero-parameter layer; nothing to do.'
 
+      type is (flatten2d_layer)
+        ! No parameters to set.
+        write(stderr, '(a)') 'Warning: calling set_params() ' &
+          // 'on a zero-parameter layer; nothing to do.'
+
       type is (reshape3d_layer)
         ! No parameters to set.
         write(stderr, '(a)') 'Warning: calling set_params() ' &

src/nf/nf_network_submodule.f90

@@ -151,7 +151,6 @@ module subroutine backward(self, output, loss)
             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)
 
@@ -283,6 +282,10 @@ module function predict_2d(self, input) result(res)
     select type(output_layer => self % layers(num_layers) % p)
       type is(dense_layer)
         res = output_layer % output
+      type is(flatten_layer)
+        res = output_layer % output
+      class default
+        error stop 'network % output not implemented for this output layer'
     end select
 
   end function predict_2d