graphnet.models.components.layers module

Class(es) implementing layers to be used in graphnet models.

class graphnet.models.components.layers.DynEdgeConv(nn, aggr, nb_neighbors, features_subset, **kwargs)[source]

Bases: EdgeConv, LightningModule

Dynamical edge convolution layer.

Construct DynEdgeConv.

Parameters:

  • nn (Callable) – The MLP/torch.Module to be used within the EdgeConv.

  • aggr (str, default: 'max') – Aggregation method to be used with EdgeConv.

  • nb_neighbors (int, default: 8) – Number of neighbours to be clustered after the EdgeConv operation.

  • features_subset (Union[Sequence[int], slice, None], default: None) – Subset of features in Data.x that should be used when dynamically performing the new graph clustering after the EdgeConv operation. Defaults to all features.

  • **kwargs (Any) – Additional features to be passed to EdgeConv.

forward(x, edge_index, batch)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • edge_index (Tensor | SparseTensor)

  • batch (Tensor | None)

class graphnet.models.components.layers.EdgeConvTito(nn, aggr, **kwargs)[source]

Bases: MessagePassing, LightningModule

Implementation of EdgeConvTito layer used in TITO solution for.

‘IceCube - Neutrinos in Deep’ kaggle competition.

Construct EdgeConvTito.

Parameters:

  • nn (Callable) – The MLP/torch.Module to be used within the EdgeConvTito.

  • aggr (str, default: 'max') – Aggregation method to be used with EdgeConvTito.

  • **kwargs (Any) – Additional features to be passed to EdgeConvTito.

reset_parameters()[source]

Reset all learnable parameters of the module.

Return type:

None

forward(x, edge_index)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor | Tuple[Tensor, Tensor])

  • edge_index (Tensor | SparseTensor)

message(x_i, x_j)[source]

Edgeconvtito message passing.

Return type:

Tensor

Parameters:

  • x_i (Tensor)

  • x_j (Tensor)

class graphnet.models.components.layers.DynTrans(layer_sizes, aggr, features_subset, n_head, **kwargs)[source]

Bases: EdgeConvTito, LightningModule

Implementation of dynTrans1 layer used in TITO solution for.

‘IceCube - Neutrinos in Deep’ kaggle competition.

Construct DynTrans.

Parameters:

  • layer_sizes (Optional[List[int]], default: None) – List of layer sizes to be used in DynTrans.

  • aggr (str, default: 'max') – Aggregation method to be used with DynTrans.

  • features_subset (Union[Sequence[int], slice, None], default: None) – Subset of features in Data.x that should be used when dynamically performing the new graph clustering after the EdgeConv operation. Defaults to all features.

  • n_head (int, default: 8) – Number of heads to be used in the multiheadattention models.

  • **kwargs (Any) – Additional features to be passed to DynTrans.

forward(x, edge_index, batch)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • edge_index (Tensor | SparseTensor)

  • batch (Tensor | None)

class graphnet.models.components.layers.DropPath(drop_prob)[source]

Bases: LightningModule

Drop paths (Stochastic Depth) per sample.

Construct DropPath.

Parameters:

drop_prob (float, default: 0.0) – Probability of dropping a path during training. If 0.0, no paths are dropped. Defaults to None.

forward(x)[source]

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

extra_repr()[source]

Return extra representation of the module.

Return type:

str

class graphnet.models.components.layers.Mlp(in_features, hidden_features, out_features, activation=<class 'torch.nn.modules.activation.GELU'>, dropout_prob)[source]

Bases: LightningModule

Multi-Layer Perceptron (MLP) module.

Construct Mlp.

Parameters:

  • in_features (int) – Number of input features.

  • hidden_features (Optional[int], default: None) – Number of hidden features. Defaults to None. If None, it is set to the value of in_features.

  • out_features (Optional[int], default: None) – Number of output features. Defaults to None. If None, it is set to the value of in_features.

  • activation (Module, default: <class 'torch.nn.modules.activation.GELU'>) – Activation layer. Defaults to nn.GELU.

  • dropout_prob (float, default: 0.0) – Dropout probability. Defaults to 0.0.

forward(x)[source]

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

class graphnet.models.components.layers.Block_rel(input_dim, num_heads, mlp_ratio, qkv_bias, qk_scale, dropout, attn_drop, drop_path, init_values, activation=<class 'torch.nn.modules.activation.GELU'>, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, attn_head_dim)[source]

Bases: LightningModule

Implementation of BEiTv2 Block.

Construct ‘Block_rel’.

Parameters:

  • input_dim (int) – Dimension of the input tensor.

  • num_heads (int) – Number of attention heads to use in the Attention_rel

  • layer.

  • mlp_ratio (float, default: 4.0) – Ratio of the hidden size of the feedforward network to the input size in the Mlp layer.

  • qkv_bias (bool, default: False) – Whether or not to include bias terms in the query, key, and value matrices in the Attention_rel layer.

  • qk_scale (Optional[float], default: None) – Scaling factor for the dot product of the query and key matrices in the Attention_rel layer.

  • dropout (float, default: 0.0) – Dropout probability to use in the Mlp layer.

  • attn_drop (float, default: 0.0) – Dropout probability to use in the Attention_rel layer.

  • drop_path (float, default: 0.0) – Probability of applying drop path regularization to the output of the layer.

  • init_values (Optional[float], default: None) – Initial value to use for the gamma_1 and gamma_2 parameters if not None.

  • activation (Module, default: <class 'torch.nn.modules.activation.GELU'>) – Activation function to use in the Mlp layer.

  • norm_layer (Module, default: <class 'torch.nn.modules.normalization.LayerNorm'>) – Normalization layer to use.

  • attn_head_dim (Optional[int], default: None) – Dimension of the attention head outputs in the Attention_rel layer.

forward(x, key_padding_mask, rel_pos_bias, kv)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • key_padding_mask (Tensor | None)

  • rel_pos_bias (Tensor | None)

  • kv (Tensor | None)

class graphnet.models.components.layers.Attention_rel(input_dim, num_heads, qkv_bias, qk_scale, attn_drop, proj_drop, attn_head_dim)[source]

Bases: LightningModule

Attention mechanism with relative position bias.

Construct ‘Attention_rel’.

Parameters:

  • input_dim (int) – Dimension of the input tensor.

  • num_heads (int, default: 8) – the number of attention heads to use (default: 8)

  • qkv_bias (bool, default: False) – whether to add bias to the query, key, and value projections. Defaults to False.

  • qk_scale (Optional[float], default: None) – a scaling factor that multiplies the dot product of query and key vectors. Defaults to None. If None, computed as :math: head_dim^(-1/2).

  • attn_drop (float, default: 0.0) – the dropout probability for the attention weights. Defaults to 0.0.

  • proj_drop (float, default: 0.0) – the dropout probability for the output of the attention module. Defaults to 0.0.

  • attn_head_dim (Optional[int], default: None) – the feature dimensionality of each attention head. Defaults to None. If None, computed as dim // num_heads.

forward(q, k, v, rel_pos_bias, key_padding_mask)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • q (Tensor)

  • k (Tensor)

  • v (Tensor)

  • rel_pos_bias (Tensor | None)

  • key_padding_mask (Tensor | None)

class graphnet.models.components.layers.Block(input_dim, num_heads, mlp_ratio, dropout, attn_drop, drop_path, init_values, activation=<class 'torch.nn.modules.activation.GELU'>, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>)[source]

Bases: LightningModule

Transformer block.

Construct ‘Block’.

Parameters:

  • input_dim (int) – Dimension of the input tensor.

  • num_heads (int) – Number of attention heads to use in the MultiheadAttention layer.

  • mlp_ratio (float, default: 4.0) – Ratio of the hidden size of the feedforward network to the input size in the Mlp layer.

  • dropout (float, default: 0.0) – Dropout probability to use in the Mlp layer.

  • attn_drop (float, default: 0.0) – Dropout probability to use in the MultiheadAttention layer.

  • drop_path (float, default: 0.0) – Probability of applying drop path regularization to the output of the layer.

  • init_values (Optional[float], default: None) – Initial value to use for the gamma_1 and gamma_2 parameters if not None.

  • activation (Module, default: <class 'torch.nn.modules.activation.GELU'>) – Activation function to use in the Mlp layer.

  • norm_layer (Module, default: <class 'torch.nn.modules.normalization.LayerNorm'>) – Normalization layer to use.

forward(x, attn_mask, key_padding_mask)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • attn_mask (Tensor | None)

  • key_padding_mask (Tensor | None)