graphnet.models.gnn.dynedge module

Implementation of the DynEdge GNN model architecture.

class graphnet.models.gnn.dynedge.DynEdge(*args, **kwargs)[source]

Bases: GNN

DynEdge (dynamical edge convolutional) model.

Construct DynEdge.

Parameters:

  • nb_inputs (int) – Number of input features on each node.

  • nb_neighbours (int, default: 8) – Number of neighbours to used in the k-nearest neighbour clustering which is performed after each (dynamical) edge convolution.

  • features_subset (Union[List[int], slice, None], default: None) – The subset of latent features on each node that are used as metric dimensions when performing the k-nearest neighbours clustering. Defaults to [0,1,2].

  • dynedge_layer_sizes (Optional[List[Tuple[int, ...]]], default: None) – The layer sizes, or latent feature dimenions, used in the DynEdgeConv layer. Each entry in dynedge_layer_sizes corresponds to a single DynEdgeConv layer; the integers in the corresponding tuple corresponds to the layer sizes in the multi-layer perceptron (MLP) that is applied within each DynEdgeConv layer. That is, a list of size-two tuples means that all DynEdgeConv layers contain a two-layer MLP. Defaults to [(128, 256), (336, 256), (336, 256), (336, 256)].

  • post_processing_layer_sizes (Optional[List[int]], default: None) – Hidden layer sizes in the MLP following the skip-concatenation of the outputs of each DynEdgeConv layer. Defaults to [336, 256].

  • readout_layer_sizes (Optional[List[int]], default: None) – Hidden layer sizes in the MLP following the post-processing _and_ optional global pooling. As this is the last layer(s) in the model, the last layer in the read-out yields the output of the DynEdge model. Defaults to [128,].

  • global_pooling_schemes (Union[List[str], str, None], default: None) – The list global pooling schemes to use. Options are: “min”, “max”, “mean”, and “sum”.

  • add_global_variables_after_pooling (bool, default: False) – Whether to add global variables after global pooling. The alternative is to added (distribute) them to the individual nodes before any convolutional operations.

  • activation_layer (Optional[str], default: None) – The activation function to use in the model.

  • add_norm_layer (bool, default: False) – Whether to add a normalization layer after each linear layer.

  • skip_readout (bool, default: False) – Whether to skip the readout layer(s). If True, the output of the last post-processing layer is returned directly.

  • args (Any)

  • kwargs (Any)

Return type:

object

forward(data)[source]

Apply learnable forward pass.

Return type:

Tensor

Parameters:

data (Data)