athena__onnx_write_submodule Submodule

Submodule containing the ONNX export procedures.

This submodule contains the routines that serialise ATHENA networks to the JSON representation used for ONNX interchange.



Functions

function resolve_onnx_export_format(format) result(ifmt)

Resolve the ONNX export format into the internal integer selector.

Arguments

Type IntentOptional Attributes Name
class(*), intent(in), optional :: format

Export format as a string name or integer selector

Return Value integer

Integer selector for the export format (1=athena_abstract, 2=onnx_expanded)


Subroutines

subroutine build_gnn_metadata(layer, prefix, metadata, num_meta)

Build the metadata entry required to reconstruct a GNN layer.

Arguments

Type IntentOptional Attributes Name
class(base_layer_type), intent(in) :: layer

GNN layer instance

character(len=*), intent(in) :: prefix

Node prefix used for this exported layer

character(len=4096), intent(inout) :: metadata(:)

Metadata strings to append to

integer, intent(inout) :: num_meta

Number of metadata entries

subroutine build_graph_inputs(network, ifmt, graph_inputs, num_inputs)

Build the ONNX graph input tensor specifications.

Arguments

Type IntentOptional Attributes Name
class(network_type), intent(in) :: network

Instance of the network

integer, intent(in) :: ifmt

Export format selector

type(onnx_tensor_type), intent(inout) :: graph_inputs(:)

Graph input tensor specifications

integer, intent(inout) :: num_inputs

Number of graph inputs

subroutine build_graph_outputs(network, ifmt, graph_outputs, num_outputs)

Build the ONNX graph output tensor specifications.

Arguments

Type IntentOptional Attributes Name
class(network_type), intent(in) :: network

Instance of the network

integer, intent(in) :: ifmt

Export format selector

type(onnx_tensor_type), intent(inout) :: graph_outputs(:)

Graph output tensor specifications

integer, intent(inout) :: num_outputs

Number of graph outputs

subroutine collect_export_nodes(network, ifmt, nodes, num_nodes, max_nodes, inits, num_inits, max_inits, gnn_metadata, num_gnn_meta)

Build the ONNX nodes, initialisers and GNN metadata.

Arguments

Type IntentOptional Attributes Name
class(network_type), intent(in) :: network

Instance of the network

integer, intent(in) :: ifmt

Export format selector

type(onnx_node_type), intent(inout) :: nodes(:)

Exported ONNX nodes

integer, intent(inout) :: num_nodes

Node counter and allocation limit

integer, intent(inout) :: max_nodes

Node counter and allocation limit

type(onnx_initialiser_type), intent(inout) :: inits(:)

Exported ONNX initialisers

integer, intent(inout) :: num_inits

Initialiser counter and allocation limit

integer, intent(inout) :: max_inits

Initialiser counter and allocation limit

character(len=4096), intent(inout) :: gnn_metadata(:)

Exported GNN metadata entries

integer, intent(inout) :: num_gnn_meta

Number of metadata entries

subroutine emit_gnn_input_renames(network, layer_id, vertex_idx, nodes, num_nodes)

Emit Identity nodes that rename GNN inputs to the expected convention.

Arguments

Type IntentOptional Attributes Name
class(network_type), intent(in) :: network

Instance of the network

integer, intent(in) :: layer_id

Layer identifier and vertex position

integer, intent(in) :: vertex_idx

Layer identifier and vertex position

type(onnx_node_type), intent(inout) :: nodes(:)

Exported ONNX nodes

integer, intent(inout) :: num_nodes

Number of exported nodes

subroutine emit_standard_node_json(network, layer_id, vertex_idx, nodes, num_nodes, max_nodes, inits, num_inits, max_inits)

Emit ONNX node records for a standard, non-GNN layer.

Arguments

Type IntentOptional Attributes Name
class(network_type), intent(in) :: network

Instance of the network

integer, intent(in) :: layer_id

Layer identifier and vertex position

integer, intent(in) :: vertex_idx

Layer identifier and vertex position

type(onnx_node_type), intent(inout) :: nodes(:)

Exported ONNX nodes

integer, intent(inout) :: num_nodes

Node counter and allocation limit

integer, intent(inout) :: max_nodes

Node counter and allocation limit

type(onnx_initialiser_type), intent(inout) :: inits(:)

Exported ONNX initialisers

integer, intent(inout) :: num_inits

Initialiser counter and allocation limit

integer, intent(inout) :: max_inits

Initialiser counter and allocation limit

subroutine initialise_export_storage(network, nodes, inits, graph_inputs, graph_outputs, gnn_metadata, max_nodes, max_inits)

Allocate the working arrays used during ONNX export.

Arguments

Type IntentOptional Attributes Name
class(network_type), intent(in) :: network

Instance of the network

type(onnx_node_type), intent(out), allocatable :: nodes(:)

Exported node storage

type(onnx_initialiser_type), intent(out), allocatable :: inits(:)

Exported initialiser storage

type(onnx_tensor_type), intent(out), allocatable :: graph_inputs(:)

Exported graph input storage

type(onnx_tensor_type), intent(out), allocatable :: graph_outputs(:)

Exported graph output storage

character(len=4096), intent(out), allocatable :: gnn_metadata(:)

Metadata storage for GNN layers

integer, intent(out) :: max_nodes

Pre-allocation sizes

integer, intent(out) :: max_inits

Pre-allocation sizes

subroutine update_pytorch_prev_output(layer, prefix, is_last_layer, output)

Resolve the downstream tensor name after emitting one PyTorch-format NOP.

Arguments

Type IntentOptional Attributes Name
class(base_layer_type), intent(in) :: layer
character(len=*), intent(in) :: prefix
logical, intent(in) :: is_last_layer
character(len=128), intent(inout) :: output

subroutine write_onnx_json_file(file, ifmt, nodes, num_nodes, inits, num_inits, graph_inputs, num_inputs, graph_outputs, num_outputs, gnn_metadata, num_gnn_meta)

Write the collected export data to disk.

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: file

Output file name

integer, intent(in) :: ifmt

Export format selector

type(onnx_node_type), intent(in) :: nodes(:)

Exported ONNX nodes

integer, intent(in) :: num_nodes

Number of exported nodes

type(onnx_initialiser_type), intent(in) :: inits(:)

Exported ONNX initialisers

integer, intent(in) :: num_inits

Number of exported initialisers

type(onnx_tensor_type), intent(in) :: graph_inputs(:)

Graph input tensor specifications

integer, intent(in) :: num_inputs

Number of graph inputs

type(onnx_tensor_type), intent(in) :: graph_outputs(:)

Graph output tensor specifications

integer, intent(in) :: num_outputs

Number of graph outputs

character(len=4096), intent(in) :: gnn_metadata(:)

GNN metadata strings

integer, intent(in) :: num_gnn_meta

Number of metadata entries


Module Subroutines

module subroutine write_onnx(file, network, format)

Export a network to ONNX JSON format.

Read more…

Arguments

Type IntentOptional Attributes Name
character(len=*), intent(in) :: file

Output file name

class(network_type), intent(in) :: network

Instance of the network

class(*), intent(in), optional :: format

Export format: 'athena_abstract' (default) or 'onnx_expanded'