init_conv Module Subroutine

  module subroutine init_conv(this, input_shape, verbose)
    !! Initialise convolutional layer
    use athena__initialiser, only: initialiser_setup
    use athena__misc_types, only: base_init_type
    implicit none

    ! Arguments
    class(conv_layer_type), intent(inout) :: this
    !! Instance of the layer
    integer, dimension(:), intent(in) :: input_shape
    !! Input shape
    integer, dimension(this%input_rank-1) :: pad_shape
    integer, optional, intent(in) :: verbose
    !! Verbosity level

    ! Local variables
    integer :: verbose_ = 0
    !! Verbosity level


    !---------------------------------------------------------------------------
    ! initialise optional arguments
    !---------------------------------------------------------------------------
    if(present(verbose)) verbose_ = verbose


    !---------------------------------------------------------------------------
    ! initialise input shape
    !---------------------------------------------------------------------------
    if(.not.allocated(this%input_shape)) call this%set_shape(input_shape)


    !---------------------------------------------------------------------------
    ! initialise padding layer, if allocated
    !---------------------------------------------------------------------------
    if(allocated(this%pad_layer))then
       call this%pad_layer%init(this%input_shape, verbose_)
       pad_shape = 2 * this%pad_layer%pad
    else
       pad_shape = 0
    end if


    !---------------------------------------------------------------------------
    ! allocate output, activation, bias, and weight shapes
    !---------------------------------------------------------------------------
    ! NOTE: INPUT SHAPE DOES NOT INCLUDE PADDING WIDTH
    ! THIS IS HANDLED AUTOMATICALLY BY THE CODE
    ! ... provide the initial input data shape and let us deal with the padding
    this%num_channels = this%input_shape(this%input_rank)
    if(allocated(this%output_shape)) deallocate(this%output_shape)
    allocate( this%output_shape(this%input_rank) )
    this%output_shape(this%input_rank) = this%num_filters
    this%output_shape(:this%input_rank-1) = floor( &
         ( &
              this%input_shape(:this%input_rank-1) + pad_shape - this%knl &
         ) / real(this%stp) &
    ) + 1
    this%num_params = this%get_num_params()
    allocate(this%weight_shape(this%input_rank + 1,1))
    this%weight_shape(:,1) = [ this%knl, this%num_channels, this%num_filters ]
    this%bias_shape = [this%num_filters]

    if(allocated(this%params)) deallocate(this%params)
    allocate(this%params(2))
    call this%params(1)%allocate([this%weight_shape(:,1), 1])
    call this%params(1)%set_requires_grad(.true.)
    this%params(1)%fix_pointer = .true.
    this%params(1)%is_sample_dependent = .false.
    this%params(1)%is_temporary = .false.
    call this%params(2)%allocate([this%bias_shape, 1])
    call this%params(2)%set_requires_grad(.true.)
    this%params(2)%fix_pointer = .true.
    this%params(2)%is_sample_dependent = .false.
    this%params(2)%is_temporary = .false.


    !---------------------------------------------------------------------------
    ! initialise weights (kernels)
    !---------------------------------------------------------------------------
    call this%kernel_init%initialise( &
         this%params(1)%val(:,1), &
         fan_in = product(this%knl)+1, fan_out = 1, &
         spacing = [ this%knl, this%num_channels, this%num_filters ] &
    )

    ! initialise biases
    !---------------------------------------------------------------------------
    call this%bias_init%initialise( &
         this%params(2)%val(:,1), &
         fan_in = product(this%knl)+1, fan_out = 1 &
    )


    !---------------------------------------------------------------------------
    ! Allocate arrays
    !---------------------------------------------------------------------------
    if(this%use_graph_input)then
       call stop_program( &
            "Graph input not supported for convolutional layer" &
       )
       return
    end if
    if(allocated(this%output)) deallocate(this%output)
    allocate( this%output(1,1) )
    if(this%z(1)%allocated) call this%z(1)%deallocate()
    if(this%z(2)%allocated) call this%z(2)%deallocate()

  end subroutine init_conv