Accumulate the gradient in dropout % backward and flush in network % update

Author: milancurcic

src/nf/nf_dropout_layer_submodule.f90

@@ -47,7 +47,8 @@ module subroutine forward(self, input)
       call shuffle(self % mask)
 
       ! Scale factor to preserve the input sum
-      self % scale = sum(input) / sum(input * self % mask)
+      self % scale = sum(input) / sum(input * self % mask)  ! input conservative approach
+      !self % scale = 1 / (1 - self % dropout_rate)  ! reference approach
 
       ! Apply dropout mask
       self % output = input * self % mask * self % scale
@@ -68,10 +69,10 @@ pure module subroutine backward(self, input, gradient)
 
     if (self % training) then
       ! Backpropagate gradient through dropout mask
-      self % gradient = gradient * self % mask * self % scale
+      self % gradient = self % gradient + gradient * self % mask * self % scale
     else
       ! In inference mode, pass through the gradient unchanged
-      self % gradient = gradient
+      self % gradient = self % gradient + gradient
     end if
   end subroutine backward
 

src/nf/nf_network_submodule.f90

@@ -543,14 +543,6 @@ module subroutine train(self, input_data, output_data, batch_size, &
       self % loss = quadratic()
     end if
 
-    ! Set all dropout layers' training mode to true.
-    do n = 2, size(self % layers)
-      select type(this_layer => self % layers(n) % p)
-        type is(dropout_layer)
-          this_layer % training = .true.
-      end select
-    end do
-
     dataset_size = size(output_data, dim=2)
 
     epoch_loop: do n = 1, epochs
@@ -640,6 +632,8 @@ module subroutine update(self, optimizer, batch_size)
         type is(conv2d_layer)
           this_layer % dw = 0
           this_layer % db = 0
+        type is(dropout_layer)
+          this_layer % gradient = 0
       end select
     end do