Update optimizer flow for layernorm

Author: milancurcic

src/nf/nf_layernorm.f90

@@ -38,7 +38,9 @@ module nf_layernorm_layer
     procedure :: init
     procedure :: get_num_params
     procedure :: get_params
+    procedure :: get_params_ptr
     procedure :: get_gradients
+    procedure :: get_gradients_ptr
     procedure :: set_params
   end type layernorm_layer
 
@@ -78,12 +80,24 @@ module function get_params(self) result(params)
     end function get_params
 
 
+    module subroutine get_params_ptr(self, g_ptr, b_ptr)
+      class(layernorm_layer), intent(in), target :: self
+      real, pointer, intent(out) :: g_ptr(:), b_ptr(:)
+    end subroutine get_params_ptr
+
+
     module function get_gradients(self) result(gradients)
       class(layernorm_layer), intent(in), target :: self
       real, allocatable :: gradients(:)
     end function get_gradients
 
 
+    module subroutine get_gradients_ptr(self, dg_ptr, db_ptr)
+      class(layernorm_layer), intent(in), target :: self
+      real, pointer, intent(out) :: dg_ptr(:), db_ptr(:)
+    end subroutine get_gradients_ptr
+
+
     module subroutine set_params(self, params)
       class(layernorm_layer), intent(in out) :: self
       real, intent(in), target :: params(:)

src/nf/nf_layernorm_submodule.f90

@@ -112,25 +112,31 @@ end function get_num_params
   module function get_params(self) result(params)
     class(layernorm_layer), intent(in), target :: self
     real, allocatable :: params(:)
+    params = [self % gamma, self % beta]
+  end function get_params
 
-    params = [ &
-      self % gamma, &
-      self % beta &
-    ]
 
-  end function get_params
+  module subroutine get_params_ptr(self, g_ptr, b_ptr)
+    class(layernorm_layer), intent(in), target :: self
+    real, pointer, intent(out) :: g_ptr(:), b_ptr(:)
+    g_ptr => self % gamma
+    b_ptr => self % beta
+  end subroutine get_params_ptr
 
 
   module function get_gradients(self) result(gradients)
     class(layernorm_layer), intent(in), target :: self
     real, allocatable :: gradients(:)
+    gradients = [self % d_gamma, self % d_beta]
+  end function get_gradients
 
-    gradients = [ &
-      self % d_gamma, &
-      self % d_beta &
-    ]
 
-  end function get_gradients
+  module subroutine get_gradients_ptr(self, dg_ptr, db_ptr)
+    class(layernorm_layer), intent(in), target :: self
+    real, pointer, intent(out) :: dg_ptr(:), db_ptr(:)
+    dg_ptr => self % d_gamma
+    db_ptr => self % d_beta
+  end subroutine get_gradients_ptr
 
 
   module subroutine set_params(self, params)

src/nf/nf_network_submodule.f90

@@ -750,6 +750,20 @@ module subroutine update(self, optimizer, batch_size)
           call self % layers(n) % optimizer % minimize(biases, db / batch_size_)
           this_layer % dw = 0
           this_layer % db = 0
+        type is(linear2d_layer)
+          call this_layer % get_params_ptr(weights, biases)
+          call this_layer % get_gradients_ptr(dw, db)
+          call self % layers(n) % optimizer % minimize(weights, dw / batch_size_)
+          call self % layers(n) % optimizer % minimize(biases, db / batch_size_)
+          this_layer % dw = 0
+          this_layer % db = 0
+        type is(layernorm_layer)
+          call this_layer % get_params_ptr(weights, biases)
+          call this_layer % get_gradients_ptr(dw, db)
+          call self % layers(n) % optimizer % minimize(weights, dw / batch_size_)
+          call self % layers(n) % optimizer % minimize(biases, db / batch_size_)
+          this_layer % d_gamma = 0
+          this_layer % d_beta = 0
       end select
     end do
 

test/test_layernorm.f90

@@ -27,14 +27,14 @@ program test_layernorm_instance
   end if
 
 contains
-  function allclose(x, y) result(res)
-    real, intent(in) :: x(:)
-    real, intent(in) :: y(:)
-    logical :: res
 
-    res = all(abs(x - y) <= (1e-06 + 1e-05 * abs(y)))
+  logical function allclose(x, y) result(res)
+    real, intent(in) :: x(:), y(:)
+    !res = all(abs(x - y) <= (1e-06 + 1e-05 * abs(y)))
+    res = all(abs(x - y) <= 1e-05)
   end function allclose
 
+
   subroutine test_layernorm_forward(layernorm_instance, input, ok)
     type(layernorm_layer), intent(in out) :: layernorm_instance
     real, intent(in out) :: input(:, :)
@@ -61,6 +61,7 @@ subroutine test_layernorm_forward(layernorm_instance, input, ok)
     end if
   end subroutine test_layernorm_forward
 
+
   subroutine test_layernorm_backward(layernorm_instance, input, gradient, ok)
     type(layernorm_layer), intent(in out) :: layernorm_instance
     real, intent(in out) :: input(:, :)
@@ -103,6 +104,7 @@ subroutine test_layernorm_backward(layernorm_instance, input, gradient, ok)
     end if
   end subroutine test_layernorm_backward
 
+
   subroutine test_layernorm_gradients(input, gradient, ok)
     real, intent(in out) :: input(:, :)
     real, intent(in out) :: gradient(:, :)
@@ -152,6 +154,7 @@ subroutine test_layernorm_gradients(input, gradient, ok)
     end if
   end subroutine test_layernorm_gradients
 
+
   subroutine test_layernorm_integration(ok)
     logical, intent(in out) :: ok
 
@@ -160,13 +163,13 @@ subroutine test_layernorm_integration(ok)
     real :: y(6) = [0.7, 0.2, 0.1, 0.1, 0.01, 0.9]
     real :: tolerance = 0.1
     integer :: epoch
-    integer :: epochs = 10000
+    integer, parameter :: num_epochs = 100000
 
-    net = network([&
-        input(2, 3),&
-        linear2d(3),&
-        layernorm(),&
-        flatten()&
+    net = network([ &
+        input(2, 3), &
+        linear2d(3), &
+        layernorm(), &
+        flatten() &
     ])
 
     ! Kaiming weights to achieve semblance of convergance
@@ -177,17 +180,18 @@ subroutine test_layernorm_integration(ok)
       l % biases = 0.2
     end select
 
-    do epoch = 1, epochs
+    do epoch = 1, num_epochs
       call net % forward(x)
       call net % backward(y)
       call net % update(optimizer=sgd(learning_rate=0.001))
       if (all(abs(net % predict(x) - y) < tolerance)) exit
     end do
 
-    if (.not. epoch <= epochs) then
+    if (.not. epoch <= num_epochs) then
       write(stderr, '(a)') &
         'linear2d + layernorm should converge in simple training.. failed'
       ok = .false.
     end if
   end subroutine test_layernorm_integration
+
 end program test_layernorm_instance