Apply optimizer to model weights without data copy (#222)

* WIP optimizer refactor w/ pointers

* WIP optimizer optimization

* Send the data to optimizer without a copy works for dense layers

* Get weights and weight gradients as 1d

* get_params_ptr and get_gradients_ptr for conv1d, conv2d, and locally_connected1d

* Define optimizer instance per layer to preserve memory across layers

* Initialization of network-wide optimizer no longer needed now that we switched to per-layer optimizer instances

* Bookkeeping for velocity, rms_gradient, etc.; optimizer tests now pass

* Update optimizer flow for linear2d

* Update optimizer flow for layernorm

* Previous bookkeeping for successive calls to optim % minimize() assumed 2 calls per batch; this is now generalized to allow any number of calls until size(params) is exhausted

* Remove get_gradients from network, layer, dense, conv1d, conv2d

* Remove optimizer as component to the network class

Author: milancurcic

src/nf/nf_conv1d_layer.f90

@@ -31,9 +31,10 @@ module nf_conv1d_layer
 
       procedure :: forward
       procedure :: backward
-      procedure :: get_gradients
+      procedure :: get_gradients_ptr
       procedure :: get_num_params
       procedure :: get_params
+      procedure :: get_params_ptr
       procedure :: init
       procedure :: set_params
 
@@ -97,14 +98,25 @@ module function get_params(self) result(params)
           !! Parameters to get
       end function get_params
 
-      module function get_gradients(self) result(gradients)
-        !! Return the gradients of this layer.
-        !! The gradients are ordered as weights first, biases second.
+      module subroutine get_params_ptr(self, w_ptr, b_ptr)
+        !! Return pointers to the parameters (weights and biases) of this layer.
         class(conv1d_layer), intent(in), target :: self
           !! A `conv1d_layer` instance
-        real, allocatable :: gradients(:)
-          !! Gradients to get
-      end function get_gradients
+        real, pointer, intent(out) :: w_ptr(:)
+          !! Pointer to the kernel weights (flattened)
+        real, pointer, intent(out) :: b_ptr(:)
+          !! Pointer to the biases
+      end subroutine get_params_ptr
+
+      module subroutine get_gradients_ptr(self, dw_ptr, db_ptr)
+        !! Return pointers to the gradients of this layer.
+        class(conv1d_layer), intent(in), target :: self
+          !! A `conv1d_layer` instance
+        real, pointer, intent(out) :: dw_ptr(:)
+          !! Pointer to the kernel weight gradients (flattened)
+        real, pointer, intent(out) :: db_ptr(:)
+          !! Pointer to the bias gradients
+      end subroutine get_gradients_ptr
 
       module subroutine set_params(self, params)
         !! Set the parameters of the layer.

src/nf/nf_conv1d_layer_submodule.f90

@@ -152,13 +152,21 @@ module function get_params(self) result(params)
     params = [ w_, self % biases]
   end function get_params
 
-  module function get_gradients(self) result(gradients)
+  module subroutine get_params_ptr(self, w_ptr, b_ptr)
     class(conv1d_layer), intent(in), target :: self
-    real, allocatable :: gradients(:)
-    real, pointer :: dw_(:) => null()
-    dw_(1:size(self % dw)) => self % dw
-    gradients = [ dw_, self % db ]
-  end function get_gradients
+    real, pointer, intent(out) :: w_ptr(:)
+    real, pointer, intent(out) :: b_ptr(:)
+    w_ptr(1:size(self % kernel)) => self % kernel
+    b_ptr => self % biases
+  end subroutine get_params_ptr
+
+  module subroutine get_gradients_ptr(self, dw_ptr, db_ptr)
+    class(conv1d_layer), intent(in), target :: self
+    real, pointer, intent(out) :: dw_ptr(:)
+    real, pointer, intent(out) :: db_ptr(:)
+    dw_ptr(1:size(self % dw)) => self % dw
+    db_ptr => self % db
+  end subroutine get_gradients_ptr
 
   module subroutine set_params(self, params)
     class(conv1d_layer), intent(in out) :: self

src/nf/nf_conv2d_layer.f90

@@ -32,9 +32,10 @@ module nf_conv2d_layer
 
     procedure :: forward
     procedure :: backward
-    procedure :: get_gradients
+    procedure :: get_gradients_ptr
     procedure :: get_num_params
     procedure :: get_params
+    procedure :: get_params_ptr
     procedure :: init
     procedure :: set_params
 
@@ -98,14 +99,25 @@ module function get_params(self) result(params)
         !! Parameters to get
     end function get_params
 
-    module function get_gradients(self) result(gradients)
-      !! Return the gradients of this layer.
-      !! The gradients are ordered as weights first, biases second.
+    module subroutine get_params_ptr(self, w_ptr, b_ptr)
+      !! Return pointers to the parameters (weights and biases) of this layer.
       class(conv2d_layer), intent(in), target :: self
         !! A `conv2d_layer` instance
-      real, allocatable :: gradients(:)
-        !! Gradients to get
-    end function get_gradients
+      real, pointer, intent(out) :: w_ptr(:)
+        !! Pointer to the kernel weights (flattened)
+      real, pointer, intent(out) :: b_ptr(:)
+        !! Pointer to the biases
+    end subroutine get_params_ptr
+
+    module subroutine get_gradients_ptr(self, dw_ptr, db_ptr)
+      !! Return pointers to the gradients of this layer.
+      class(conv2d_layer), intent(in), target :: self
+        !! A `conv2d_layer` instance
+      real, pointer, intent(out) :: dw_ptr(:)
+        !! Pointer to the kernel weight gradients (flattened)
+      real, pointer, intent(out) :: db_ptr(:)
+        !! Pointer to the bias gradients
+    end subroutine get_gradients_ptr
 
     module subroutine set_params(self, params)
       !! Set the parameters of the layer.

src/nf/nf_conv2d_layer_submodule.f90

@@ -204,21 +204,23 @@ module function get_params(self) result(params)
 
   end function get_params
 
-
-  module function get_gradients(self) result(gradients)
+  
+  module subroutine get_params_ptr(self, w_ptr, b_ptr)
     class(conv2d_layer), intent(in), target :: self
-    real, allocatable :: gradients(:)
-
-    real, pointer :: dw_(:) => null()
+    real, pointer, intent(out) :: w_ptr(:)
+    real, pointer, intent(out) :: b_ptr(:)
+    w_ptr(1:size(self % kernel)) => self % kernel
+    b_ptr => self % biases
+  end subroutine get_params_ptr
 
-    dw_(1:size(self % dw)) => self % dw
 
-    gradients = [ &
-      dw_, &
-      self % db &
-    ]
-
-  end function get_gradients
+  module subroutine get_gradients_ptr(self, dw_ptr, db_ptr)
+    class(conv2d_layer), intent(in), target :: self
+    real, pointer, intent(out) :: dw_ptr(:)
+    real, pointer, intent(out) :: db_ptr(:)
+    dw_ptr(1:size(self % dw)) => self % dw
+    db_ptr => self % db
+  end subroutine get_gradients_ptr
 
 
   module subroutine set_params(self, params)

src/nf/nf_dense_layer.f90

@@ -33,9 +33,10 @@ module nf_dense_layer
 
     procedure :: backward
     procedure :: forward
-    procedure :: get_gradients
+    procedure :: get_gradients_ptr
     procedure :: get_num_params
     procedure :: get_params
+    procedure :: get_params_ptr
     procedure :: init
     procedure :: set_params
 
@@ -96,14 +97,17 @@ module function get_params(self) result(params)
         !! Parameters of this layer
     end function get_params
 
-    module function get_gradients(self) result(gradients)
-      !! Return the gradients of this layer.
-      !! The gradients are ordered as weights first, biases second.
+    module subroutine get_params_ptr(self, w_ptr, b_ptr)
       class(dense_layer), intent(in), target :: self
-        !! Dense layer instance
-      real, allocatable :: gradients(:)
-        !! Gradients of this layer
-    end function get_gradients
+      real, pointer, intent(out) :: w_ptr(:)
+      real, pointer, intent(out) :: b_ptr(:)
+    end subroutine get_params_ptr
+
+    module subroutine get_gradients_ptr(self, dw_ptr, db_ptr)
+      class(dense_layer), intent(in), target :: self
+      real, pointer, intent(out) :: dw_ptr(:)
+      real, pointer, intent(out) :: db_ptr(:)
+    end subroutine get_gradients_ptr
 
     module subroutine set_params(self, params)
       !! Set the parameters of this layer.

src/nf/nf_dense_layer_submodule.f90

@@ -77,20 +77,22 @@ module function get_params(self) result(params)
   end function get_params
 
 
-  module function get_gradients(self) result(gradients)
+  module subroutine get_params_ptr(self, w_ptr, b_ptr)
     class(dense_layer), intent(in), target :: self
-    real, allocatable :: gradients(:)
+    real, pointer, intent(out) :: w_ptr(:)
+    real, pointer, intent(out) :: b_ptr(:)
+    w_ptr(1:size(self % weights)) => self % weights
+    b_ptr => self % biases
+  end subroutine get_params_ptr
 
-    real, pointer :: dw_(:) => null()
 
-    dw_(1:size(self % dw)) => self % dw
-
-    gradients = [ &
-      dw_, &
-      self % db &
-    ]
-
-  end function get_gradients
+  module subroutine get_gradients_ptr(self, dw_ptr, db_ptr)
+    class(dense_layer), intent(in), target :: self
+    real, pointer, intent(out) :: dw_ptr(:)
+    real, pointer, intent(out) :: db_ptr(:)
+    dw_ptr(1:size(self % dw)) => self % dw
+    db_ptr => self % db
+  end subroutine get_gradients_ptr
 
 
   module subroutine set_params(self, params)

src/nf/nf_layer.f90

@@ -22,13 +22,13 @@ module nf_layer
     integer, allocatable :: layer_shape(:)
     integer, allocatable :: input_layer_shape(:)
     logical :: initialized = .false.
+    class(optimizer_base_type), allocatable :: optimizer
 
   contains
 
     procedure :: forward
     procedure :: get_num_params
     procedure :: get_params
-    procedure :: get_gradients
     procedure :: set_params
     procedure :: init
     procedure :: print_info
@@ -160,14 +160,6 @@ module function get_params(self) result(params)
         !! Parameters of this layer
     end function get_params
 
-    module function get_gradients(self) result(gradients)
-      !! Returns the gradients of this layer.
-      class(layer), intent(in) :: self
-        !! Layer instance
-      real, allocatable :: gradients(:)
-        !! Gradients of this layer
-    end function get_gradients
-
     module subroutine set_params(self, params)
       !! Returns the parameters of this layer.
       class(layer), intent(in out) :: self

src/nf/nf_layer_submodule.f90

@@ -682,50 +682,6 @@ module function get_params(self) result(params)
 
   end function get_params
 
-  module function get_gradients(self) result(gradients)
-    class(layer), intent(in) :: self
-    real, allocatable :: gradients(:)
-
-    select type (this_layer => self % p)
-      type is (input1d_layer)
-        ! No gradients to get.
-      type is (input2d_layer)
-        ! No gradients to get.
-      type is (input3d_layer)
-        ! No gradients to get.
-      type is (dense_layer)
-        gradients = this_layer % get_gradients()
-      type is (dropout_layer)
-        ! No gradients to get.
-      type is (conv1d_layer)
-        gradients = this_layer % get_gradients()
-      type is (conv2d_layer)
-        gradients = this_layer % get_gradients()
-      type is (locally_connected1d_layer)
-        gradients = this_layer % get_gradients()
-      type is (maxpool1d_layer)
-        ! No gradients to get.
-      type is (maxpool2d_layer)
-        ! No gradients to get.
-      type is (flatten_layer)
-        ! No gradients to get.
-      type is (reshape2d_layer)
-        ! No parameters to get.
-      type is (reshape3d_layer)
-        ! No gradients to get.
-      type is (linear2d_layer)
-        gradients = this_layer % get_gradients()
-      type is (self_attention_layer)
-        gradients = this_layer % get_gradients()
-      type is (embedding_layer)
-        gradients = this_layer % get_gradients()
-      type is (layernorm_layer)
-        gradients = this_layer % get_gradients()
-      class default
-        error stop 'Unknown layer type.'
-    end select
-
-  end function get_gradients
 
   module subroutine set_params(self, params)
     class(layer), intent(in out) :: self

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)