attention_blocks

Attention and transformer blocks used in graphnet models.

class graphnet.models.components.attention_blocks.DropPath(drop_prob)

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)

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

extra_repr()

Return extra representation of the module.

Return type:

str

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

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)

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

class graphnet.models.components.attention_blocks.Attention_rel(input_dim, num_heads, qkv_bias, qk_scale, attn_drop, proj_drop, attn_head_dim)

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)

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.attention_blocks.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)

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)

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.attention_blocks.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'>)

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)

Forward pass.

Return type:

Tensor

Parameters:

• x (Tensor)

• attn_mask (Tensor | None)

• key_padding_mask (Tensor | None)