module athena__base_layer !! Module containing the abstract base layer type !! !! This module contains the abstract base layer type, from which all other !! layers are derived. The module also contains the abstract derived types !! for the following layer types: !! - padding !! - pooling !! - dropout !! - learnable !! - convolutional !! - batch normalisation !! !! The following procedures are based on code from the neural-fortran library !! https://github.com/modern-fortran/neural-fortran/blob/main/src/nf/nf_layer.f90 use coreutils, only: real32 use athena__clipper, only: clip_type use athena__misc_types, only: base_actv_type, base_init_type, facets_type, & onnx_attribute_type, onnx_node_type, onnx_initialiser_type, & onnx_tensor_type use diffstruc, only: array_type use athena__diffstruc_extd, only: array_ptr_type use graphstruc, only: graph_type implicit none private public :: base_layer_type public :: pad_layer_type public :: pool_layer_type public :: drop_layer_type public :: learnable_layer_type public :: conv_layer_type public :: batch_layer_type public :: merge_layer_type !------------------------------------------------------------------------------- ! layer abstract type !------------------------------------------------------------------------------- type, abstract :: base_layer_type !! Type for base layer, from which all other layers are derived integer :: id !! Unique identifier integer :: input_rank = 0 !! Rank of input data integer :: output_rank = 0 !! Rank of output data logical :: inference = .false. !! Inference mode logical :: use_graph_input = .false. !! Use graph input logical :: use_graph_output = .false. !! Use graph output character(:), allocatable :: name !! Layer name character(4) :: type = 'base' !! Layer type character(20) :: subtype = repeat(" ",20) type(graph_type), allocatable, dimension(:) :: graph !! Graph structure of input data class(array_type), allocatable, dimension(:,:) :: output !! Output integer, allocatable, dimension(:) :: input_shape !! Input shape integer, allocatable, dimension(:) :: output_shape !! Output shape contains procedure, pass(this) :: set_rank => set_rank_base !! Set the input and output ranks of the layer procedure, pass(this) :: set_shape => set_shape_base !! Set the input shape of the layer procedure, pass(this) :: get_num_params => get_num_params_base !! Get the number of parameters in the layer procedure, pass(this) :: print => print_base !! Print the layer to a file with additional information procedure, pass(this) :: print_to_unit => print_to_unit_base !! Print the layer to a unit procedure, pass(this) :: get_attributes => get_attributes_base !! Get the attributes of the layer (for ONNX export) procedure, pass(this) :: extract_output => extract_output_base !! Extract the output of the layer as a standard real array procedure(initialise), deferred, pass(this) :: init !! Initialise the layer procedure, pass(this) :: forward => forward_base !! Forward pass of layer procedure, pass(this) :: forward_eval => forward_eval_base !! Forward pass of layer and return output for evaluation procedure, pass(this) :: nullify_graph => nullify_graph_base !! Nullify the forward pass data of the layer to free memory !! Forward pass of layer using derived array_type procedure(read_layer), deferred, pass(this) :: read !! Read layer from file procedure, pass(this) :: build_from_onnx => build_from_onnx_base !! Build layer from ONNX node and initialiser procedure, pass(this) :: set_graph => set_graph_base !! Set the graph structure of the input data !! this is adjacency and edge weighting procedure, pass(this) :: emit_onnx_nodes => emit_onnx_nodes_base !! Emit ONNX JSON nodes for this layer (format-aware and polymorphic) procedure, pass(this) :: emit_onnx_graph_inputs => & emit_onnx_graph_inputs_base !! Emit graph input tensor declarations for this layer end type base_layer_type interface module subroutine print_base(this, file, unit, print_header_footer) !! Print the layer to a file with additional information class(base_layer_type), intent(in) :: this !! Instance of the layer character(*), optional, intent(in) :: file !! File name integer, optional, intent(in) :: unit !! Unit number logical, optional, intent(in) :: print_header_footer !! Boolean whether to print header and footer end subroutine print_base module subroutine print_to_unit_base(this, unit) !! Print the layer to a file class(base_layer_type), intent(in) :: this !! Instance of the layer integer, intent(in) :: unit !! File unit end subroutine print_to_unit_base module function get_attributes_base(this) result(attributes) !! Get the attributes of the layer (for ONNX export) class(base_layer_type), intent(in) :: this !! Instance of the layer type(onnx_attribute_type), allocatable, dimension(:) :: attributes !! Attributes of the layer end function get_attributes_base module subroutine set_rank_base(this, input_rank, output_rank) !! Set the input and output ranks of the layer class(base_layer_type), intent(inout) :: this !! Instance of the layer integer, intent(in) :: input_rank !! Input rank integer, intent(in) :: output_rank !! Output rank end subroutine set_rank_base module subroutine set_shape_base(this, input_shape) !! Set the input shape of the layer class(base_layer_type), intent(inout) :: this !! Instance of the layer integer, dimension(:), intent(in) :: input_shape !! Input shape end subroutine set_shape_base module subroutine extract_output_base(this, output) !! Extract the output of the layer as a standard real array class(base_layer_type), intent(in) :: this !! Instance of the layer real(real32), dimension(..), allocatable, intent(out) :: output !! Output values end subroutine extract_output_base pure module function get_num_params_base(this) result(num_params) class(base_layer_type), intent(in) :: this integer :: num_params end function get_num_params_base end interface interface module subroutine initialise(this, input_shape, verbose) !! Initialise the layer class(base_layer_type), intent(inout) :: this !! Instance of the layer integer, dimension(:), intent(in) :: input_shape !! Input shape integer, optional, intent(in) :: verbose !! Verbosity level end subroutine initialise end interface interface pure module function get_num_params(this) result(num_params) !! Get number of parameters in layer class(base_layer_type), intent(in) :: this !! Instance of the layer integer :: num_params !! Number of parameters end function get_num_params end interface interface module subroutine forward_base(this, input) !! Forward pass of layer class(base_layer_type), intent(inout) :: this !! Instance of the layer class(array_type), dimension(:,:), intent(in) :: input !! Input data end subroutine forward_base module function forward_eval_base(this, input) result(output) !! Forward pass of layer and return output for evaluation class(base_layer_type), intent(inout), target :: this !! Instance of the layer class(array_type), dimension(:,:), intent(in) :: input !! Input data type(array_type), pointer :: output(:,:) !! Output data end function forward_eval_base module subroutine set_graph_base(this, graph) !! Set the graph structure of the input data class(base_layer_type), intent(inout) :: this !! Instance of the layer type(graph_type), dimension(:), intent(in) :: graph !! Graph structure of input data end subroutine set_graph_base module subroutine nullify_graph_base(this) !! Nullify the forward pass data of the layer to free memory class(base_layer_type), intent(inout) :: this !! Instance of the layer end subroutine nullify_graph_base end interface interface module subroutine read_layer(this, unit, verbose) !! Read layer from file class(base_layer_type), intent(inout) :: this !! Instance of the layer integer, intent(in) :: unit !! File unit integer, optional, intent(in) :: verbose !! Verbosity level end subroutine read_layer module subroutine build_from_onnx_base( & this, node, initialisers, value_info, verbose & ) !! Build layer from ONNX node class(base_layer_type), intent(inout) :: this !! Instance of the layer type(onnx_node_type), intent(in) :: node !! ONNX node type(onnx_initialiser_type), dimension(:), intent(in) :: initialisers !! ONNX initialisers type(onnx_tensor_type), dimension(:), intent(in) :: value_info !! ONNX value info integer, intent(in) :: verbose !! Verbosity level end subroutine build_from_onnx_base module subroutine emit_onnx_nodes_base( & this, prefix, & nodes, num_nodes, max_nodes, & inits, num_inits, max_inits, & input_name, is_last_layer, format & ) !! Emit ONNX JSON nodes for this layer !! Default implementation does nothing; override for GNN/NOP layers class(base_layer_type), intent(in) :: this !! Instance of the layer character(*), intent(in) :: prefix !! Node name prefix (e.g. "node_2") type(onnx_node_type), intent(inout), dimension(:) :: nodes !! Accumulator for ONNX nodes integer, intent(inout) :: num_nodes !! Current number of nodes integer, intent(in) :: max_nodes !! Maximum capacity type(onnx_initialiser_type), intent(inout), dimension(:) :: inits !! Accumulator for ONNX initialisers integer, intent(inout) :: num_inits !! Current number of initialisers integer, intent(in) :: max_inits !! Maximum capacity character(*), optional, intent(in) :: input_name !! Upstream tensor name used by sequential expanded ONNX format logical, optional, intent(in) :: is_last_layer !! Whether this is the last non-input layer in the network integer, optional, intent(in) :: format !! Export format selector !! 1 = ONNX athena abstract format (default) !! 2 = ONNX expanded format end subroutine emit_onnx_nodes_base module subroutine emit_onnx_graph_inputs_base( & this, prefix, & graph_inputs, num_inputs & ) !! Emit graph input tensor declarations for this layer !! Default implementation does nothing; override for GNN layers class(base_layer_type), intent(in) :: this !! Instance of the layer character(*), intent(in) :: prefix !! Input name prefix (e.g. "input_1") type(onnx_tensor_type), intent(inout), dimension(:) :: graph_inputs !! Accumulator for graph inputs integer, intent(inout) :: num_inputs !! Current number of inputs end subroutine emit_onnx_graph_inputs_base end interface type, abstract, extends(base_layer_type) :: pad_layer_type !! Type for padding layers integer :: num_channels !! Number of channels integer :: imethod = 0 !! Method for padding integer, allocatable, dimension(:) :: pad !! Padding size character(len=20) :: method = 'valid' !! Padding method integer, allocatable, dimension(:,:) :: orig_bound, dest_bound !! Original and destination bounds type(facets_type), dimension(:), allocatable :: facets !! Facets of the layer contains procedure, pass(this) :: init => init_pad !! Initialise the layer procedure, pass(this) :: print_to_unit => print_to_unit_pad !! Print layer to unit end type pad_layer_type interface module subroutine print_to_unit_pad(this, unit) !! Print layer to unit class(pad_layer_type), intent(in) :: this !! Instance of the layer integer, intent(in) :: unit !! File unit end subroutine print_to_unit_pad module subroutine init_pad(this, input_shape, verbose) class(pad_layer_type), intent(inout) :: this integer, dimension(:), intent(in) :: input_shape integer, optional, intent(in) :: verbose end subroutine init_pad end interface type, abstract, extends(base_layer_type) :: pool_layer_type !! Type for pooling layers integer, allocatable, dimension(:) :: pool, strd !! Pooling and stride sizes integer :: num_channels !! Number of channels class(pad_layer_type), allocatable :: pad_layer !! Padding layer contains procedure, pass(this) :: init => init_pool !! Initialise the layer procedure, pass(this) :: print_to_unit => print_to_unit_pool !! Print layer to unit procedure, pass(this) :: get_attributes => get_attributes_pool !! Get the attributes of the layer (for ONNX export) end type pool_layer_type interface module subroutine print_to_unit_pool(this, unit) !! Print layer to unit class(pool_layer_type), intent(in) :: this !! Instance of the layer integer, intent(in) :: unit !! File unit end subroutine print_to_unit_pool module subroutine init_pool(this, input_shape, verbose) class(pool_layer_type), intent(inout) :: this integer, dimension(:), intent(in) :: input_shape integer, optional, intent(in) :: verbose end subroutine init_pool module function get_attributes_pool(this) result(attributes) !! Get the attributes of the layer (for ONNX export) class(pool_layer_type), intent(in) :: this !! Instance of the layer type(onnx_attribute_type), allocatable, dimension(:) :: attributes !! Attributes of the layer end function get_attributes_pool end interface type, abstract, extends(base_layer_type) :: drop_layer_type !! Type for dropout layers real(real32) :: rate = 0.1_real32 !! Dropout rate, rate = 1 - keep_prob -- typical = 0.05-0.25 contains procedure(generate_mask), deferred, pass(this) :: generate_mask !! Generate dropout mask end type drop_layer_type abstract interface subroutine generate_mask(this) !! Generate dropout mask import :: drop_layer_type class(drop_layer_type), intent(inout) :: this !! Instance of the layer end subroutine generate_mask end interface type, abstract, extends(base_layer_type) :: merge_layer_type !! Type for merge layers (i.e. add, multiply, concatenate) integer :: merge_mode = 1 !! Integer code for fundamental merge method !! 1 = pointwise !! 2 = concatenate !! 3 = reduction !! 4 = parametric (NOT IMPLEMENTED) character(len=20) :: method !! Merge method integer :: num_input_layers = 0 !! Number of input layers integer, allocatable, dimension(:) :: input_layer_ids !! IDs of input layers contains procedure(combine_merge), deferred, pass(this) :: combine !! Merge two layers (forward) procedure(calc_input_shape), deferred, pass(this) :: calc_input_shape !! Calculate input shape based on shapes of input layers end type merge_layer_type interface module subroutine combine_merge(this, input_list) !! Combine two layers (forward) class(merge_layer_type), intent(inout) :: this !! Instance of the layer type(array_ptr_type), dimension(:), intent(in) :: input_list !! Input values end subroutine combine_merge module function calc_input_shape(this, input_shapes) result(input_shape) !! Calculate input shape based on shapes of input layers class(merge_layer_type), intent(in) :: this !! Instance of the layer integer, dimension(:,:), intent(in) :: input_shapes !! Input shapes integer, allocatable, dimension(:) :: input_shape !! Calculated input shape end function calc_input_shape end interface type, abstract, extends(base_layer_type) :: learnable_layer_type !! Type for layers with learnable parameters integer :: num_params = 0 !! Number of learnable parameters logical :: use_bias = .false. !! Layer has bias integer, allocatable, dimension(:,:) :: weight_shape !! Shape of weights integer, allocatable, dimension(:) :: bias_shape !! Shape of biases type(array_type), allocatable, dimension(:) :: params !! Learnable parameters character(len=14) :: kernel_initialiser='', bias_initialiser='' !! Initialisers for kernel and bias class(base_init_type), allocatable :: kernel_init, bias_init !! Initialisers for kernel and bias class(base_actv_type), allocatable :: activation !! Activation function contains procedure, pass(this) :: get_params => get_params !! Get learnable parameters of layer procedure, pass(this) :: set_params => set_params !! Set learnable parameters of layer procedure, pass(this) :: get_gradients => get_gradients !! Get parameter gradients of layer procedure, pass(this) :: set_gradients => set_gradients !! Set learnable parameters of layer procedure, pass(this) :: reduce => reduce_learnable !! Merge another learnable layer into this one procedure :: add_t_t => add_learnable !! Add two layers generic :: operator(+) => add_t_t !! Operator overloading for addition end type learnable_layer_type interface module subroutine reduce_learnable(this, input) !! Merge another learnable layer into this one class(learnable_layer_type), intent(inout) :: this !! Instance of the layer class(learnable_layer_type), intent(in) :: input !! Other layer to merge end subroutine reduce_learnable module function add_learnable(a, b) result(output) !! Add two layers class(learnable_layer_type), intent(in) :: a, b !! Instances of the layers class(learnable_layer_type), allocatable :: output !! Output layer end function add_learnable end interface interface pure module function get_params(this) result(params) !! Get learnable parameters of layer class(learnable_layer_type), intent(in) :: this !! Instance of the layer real(real32), dimension(this%num_params) :: params !! Learnable parameters end function get_params module subroutine set_params(this, params) !! Set learnable parameters of layer class(learnable_layer_type), intent(inout) :: this !! Instance of the layer real(real32), dimension(this%num_params), intent(in) :: params !! Learnable parameters end subroutine set_params pure module function get_gradients(this, clip_method) result(gradients) !! Get parameter gradients of layer class(learnable_layer_type), intent(in) :: this !! Instance of the layer type(clip_type), optional, intent(in) :: clip_method !! Clip method real(real32), dimension(this%num_params) :: gradients !! Parameter gradients end function get_gradients module subroutine set_gradients(this, gradients) !! Set learnable parameters of layer class(learnable_layer_type), intent(inout) :: this !! Instance of the layer real(real32), dimension(..), intent(in) :: gradients !! Learnable parameters end subroutine set_gradients end interface type, abstract, extends(learnable_layer_type) :: conv_layer_type integer :: num_channels !! Number of channels integer :: num_filters !! Number of filters integer, allocatable, dimension(:) :: knl, stp, dil !! Kernel, stride, and dilation sizes real(real32), pointer :: bias(:) => null() !! Bias pointer class(pad_layer_type), allocatable :: pad_layer !! Optional preprocess padding layer class(array_type), allocatable :: di_padded !! Padded input gradients type(array_type), dimension(2) :: z !! Temporary arrays for forward propagation contains procedure, pass(this) :: get_num_params => get_num_params_conv !! Get the number of parameters in the layer procedure, pass(this) :: init => init_conv !! Initialise the layer procedure, pass(this) :: get_attributes => get_attributes_conv !! Get the attributes of the layer (for ONNX export) procedure, pass(this) :: print_to_unit => print_to_unit_conv !! Print layer to unit end type conv_layer_type interface pure module function get_num_params_conv(this) result(num_params) class(conv_layer_type), intent(in) :: this integer :: num_params end function get_num_params_conv module subroutine init_conv(this, input_shape, verbose) class(conv_layer_type), intent(inout) :: this integer, dimension(:), intent(in) :: input_shape integer, optional, intent(in) :: verbose end subroutine init_conv module function get_attributes_conv(this) result(attributes) class(conv_layer_type), intent(in) :: this type(onnx_attribute_type), allocatable, dimension(:) :: attributes end function get_attributes_conv module subroutine print_to_unit_conv(this, unit) !! Print layer to unit class(conv_layer_type), intent(in) :: this !! Instance of the layer integer, intent(in) :: unit !! File unit end subroutine print_to_unit_conv end interface type, abstract, extends(learnable_layer_type) :: batch_layer_type !! Type for batch normalisation layers integer :: num_channels !! Number of channels real(real32) :: momentum = 0.99_real32 !! Momentum factor !! NOTE: if momentum = 0, mean and variance batch-dependent values !! NOTE: if momentum > 0, mean and variance are running averages real(real32) :: epsilon = 0.001_real32 !! Epsilon factor real(real32) :: gamma_init_mean = 1._real32, gamma_init_std = 0.01_real32 !! Initialisation parameters for gamma real(real32) :: beta_init_mean = 0._real32, beta_init_std = 0.01_real32 !! Initialisation parameters for beta class(base_init_type), allocatable :: moving_mean_init, moving_variance_init !! Initialisers for moving mean and variance real(real32), allocatable, dimension(:) :: mean, variance !! Mean and variance (not learnable) contains procedure, pass(this) :: get_num_params => get_num_params_batch !! Get the number of parameters in the layer procedure, pass(this) :: init => init_batch !! Initialise the layer procedure, pass(this) :: print_to_unit => print_to_unit_batch !! Print layer to unit procedure, pass(this) :: get_attributes => get_attributes_batch !! Get the attributes of the layer (for ONNX export) end type batch_layer_type interface pure module function get_num_params_batch(this) result(num_params) class(batch_layer_type), intent(in) :: this integer :: num_params end function get_num_params_batch module subroutine init_batch(this, input_shape, verbose) class(batch_layer_type), intent(inout) :: this integer, dimension(:), intent(in) :: input_shape integer, optional, intent(in) :: verbose end subroutine init_batch module subroutine print_to_unit_batch(this, unit) !! Print layer to unit class(batch_layer_type), intent(in) :: this !! Instance of the layer integer, intent(in) :: unit !! File unit end subroutine print_to_unit_batch module function get_attributes_batch(this) result(attributes) class(batch_layer_type), intent(in) :: this !! Instance of the layer type(onnx_attribute_type), allocatable, dimension(:) :: attributes !! Attributes of the layer end function get_attributes_batch end interface end module athena__base_layer