Specific procedures for batch output

Author: milancurcic

src/nf/nf_network.f90

@@ -23,13 +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_1d
     procedure, private :: output_batch_3d
 
-    generic :: forward => forward_1d, forward_3d, forward_batch_3d
-    generic :: output => output_1d, output_3d, output_batch_3d
+    generic :: forward => forward_1d, forward_3d
+    generic :: output => output_1d, output_3d, output_batch_1d, output_batch_3d
 
   end type network
 
@@ -85,20 +85,6 @@ 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
@@ -123,14 +109,24 @@ module function output_3d(self, input) result(res)
         !! Output of the network
     end function output_3d
 
+    module function output_batch_1d(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 last dimension is the batch
+      real, allocatable :: res(:,:)
+        !! Output of the network; the last dimension is the batch
+    end function output_batch_1d
+
     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
+        !! Input data; the last dimension is the batch
       real, allocatable :: res(:,:)
-        !! Output of the network; the 2nd dimension is the batch
+        !! Output of the network; the last dimension is the batch
     end function output_batch_3d
 
   end interface output

src/nf/nf_network_submodule.f90

@@ -244,23 +244,6 @@ 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(:)
@@ -308,28 +291,66 @@ module function output_3d(self, input) result(res)
   end function output_3d
 
 
+  module function output_batch_1d(self, input) result(res)
+    class(network), intent(in out) :: self
+    real, intent(in) :: input(:,:)
+    real, allocatable :: res(:,:)
+    integer :: i, batch_size, num_layers, output_size
+
+    num_layers = size(self % layers)
+    batch_size = size(input, dim=rank(input))
+    output_size = product(self % layers(num_layers) % layer_shape)
+
+    allocate(res(output_size, batch_size))
+
+    batch: do concurrent(i = 1:size(res, dim=2))
+
+      call self % forward(input(:,i))
+
+      select type(output_layer => self % layers(num_layers) % p)
+        type is(dense_layer)
+          res(:,i) = output_layer % output
+        type is(flatten_layer)
+          res(:,i) = output_layer % output
+        class default
+          error stop 'network % output not implemented for this output layer'
+      end select
+
+    end do batch
+
+  end function output_batch_1d
+
+
   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
+    integer :: i, batch_size, num_layers, output_size
 
     num_layers = size(self % layers)
+    batch_size = size(input, dim=rank(input))
+    output_size = product(self % layers(num_layers) % layer_shape)
 
-    call self % forward(input)
+    allocate(res(output_size, batch_size))
+
+    batch: do concurrent(i = 1:batch_size)
+
+    call self % forward(input(:,:,:,i))
 
     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.)
+        res(:,i) = pack(output_layer % output, .true.)
       type is(dense_layer)
-        res = output_layer % output
+        res(:,i) = output_layer % output
       type is(flatten_layer)
-        res = output_layer % output
+        res(:,i) = output_layer % output
       class default
         error stop 'network % output not implemented for this output layer'
     end select
 
+    end do batch
+
   end function output_batch_3d