Add specific network methods for batch forward and output

Author: milancurcic

src/nf/nf_network.f90

@@ -23,11 +23,13 @@ module nf_network
 
     procedure, private :: forward_1d
     procedure, private :: forward_3d
+    procedure, private :: forward_batch_3d
     procedure, private :: output_1d
     procedure, private :: output_3d
+    procedure, private :: output_batch_3d
 
-    generic :: forward => forward_1d, forward_3d
-    generic :: output => output_1d, output_3d
+    generic :: forward => forward_1d, forward_3d, forward_batch_3d
+    generic :: output => output_1d, output_3d, output_batch_3d
 
   end type network
 
@@ -83,6 +85,20 @@ pure module subroutine forward_3d(self, input)
         !! 3-d input data
     end subroutine forward_3d
 
+    pure module subroutine forward_batch_3d(self, input)
+      !! Apply a forward pass of a batch of data through the network.
+      !!
+      !! This changes the state of layers on the network.
+      !! Typically used only internally from the `train` method,
+      !! but can be invoked by the user when creating custom optimizers.
+      !!
+      !! This specific subroutine is for 3-d input data.
+      class(network), intent(in out) :: self
+        !! Network instance
+      real, intent(in) :: input(:,:,:,:)
+        !! 3-d input data; the 4th dimension is the batch
+    end subroutine forward_batch_3d
+
   end interface forward
 
   interface output
@@ -107,6 +123,16 @@ module function output_3d(self, input) result(res)
         !! Output of the network
     end function output_3d
 
+    module function output_batch_3d(self, input) result(res)
+      !! Return the output of the network given an input batch of 3-d data.
+      class(network), intent(in out) :: self
+        !! Network instance
+      real, intent(in) :: input(:,:,:,:)
+        !! Input data; the 4th dimension is the batch
+      real, allocatable :: res(:,:)
+        !! Output of the network; the 2nd dimension is the batch
+    end function output_batch_3d
+
   end interface output
 
   interface

src/nf/nf_network_submodule.f90

@@ -244,6 +244,23 @@ pure module subroutine forward_3d(self, input)
   end subroutine forward_3d
 
 
+  pure module subroutine forward_batch_3d(self, input)
+    class(network), intent(in out) :: self
+    real, intent(in) :: input(:,:,:,:)
+    integer :: n
+
+    ! Set the input array into the input layer
+    select type(input_layer => self % layers(1) % p); type is(input3d_layer)
+      call input_layer % set(input)
+    end select
+
+    do n = 2, size(self % layers)
+      call self % layers(n) % forward(self % layers(n - 1))
+    end do
+
+  end subroutine forward_batch_3d
+
+
   module function output_1d(self, input) result(res)
     class(network), intent(in out) :: self
     real, intent(in) :: input(:)
@@ -291,6 +308,31 @@ module function output_3d(self, input) result(res)
   end function output_3d
 
 
+  module function output_batch_3d(self, input) result(res)
+    class(network), intent(in out) :: self
+    real, intent(in) :: input(:,:,:,:)
+    real, allocatable :: res(:,:)
+    integer :: num_layers
+
+    num_layers = size(self % layers)
+
+    call self % forward(input)
+
+    select type(output_layer => self % layers(num_layers) % p)
+      type is(conv2d_layer)
+        !FIXME flatten the result for now; find a better solution
+        res = pack(output_layer % output, .true.)
+      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 output_batch_3d
+
+
   module subroutine print_info(self)
     class(network), intent(in) :: self
     call self % layers % print_info()