opengt.network¶

class opengt.network.bga_model.BGA(dim_in: int, dim_out: int, dropout1=0.5, dropout2=0.1)[source]¶

Bases: Module

Bilevel Graph Attention model. Used in CoBFormer model. Adapted from https://github.com/null-xyj/CoBFormer

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
dropout1 (float) – Dropout rate for the final layer.
dropout2 (float) – Dropout rate for the BGA layers.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.patch (torch.Tensor): Patch indices.

Output:

batch.x (torch.Tensor): Output node features after applying the BGA model.

class opengt.network.big_bird.BigBird(dim_in, dim_out)[source]¶

Bases: Module

BigBird without edge features. This model disregards edge features and runs a linear transformer over a set of node features only. BirBird applies random sparse attention to the input sequence - the longer the sequence the closer it is to O(N) https://arxiv.org/abs/2007.14062

class opengt.network.cobformer.CoBFormer(dim_in: int, dim_out: int)[source]¶

Bases: Module

CoBFormer model. Only supports transductive node level tasks. Adapted from https://github.com/null-xyj/CoBFormer

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.alpha (float): Balance factor for GNN and BGA loss. - cfg.gt.tau (float): Temperature parameter for softmax. - cfg.gt.layer_type (str): Type of GNN layer to use. e.g., ‘GCN’. - cfg.gnn.layers (int): Number of GNN layers.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.
batch.patch (torch.Tensor): Patch indices.
batch.y (torch.Tensor): Input labels.
batch.split (str): Split type (train, val, test).

Output:

pred (torch.Tensor): Predicted node features after applying the CoBFormer model. true (torch.Tensor): True labels. extra_loss (torch.Tensor): Extra loss term for GNN and BGA cotraining.

class opengt.network.custom_gnn.CustomGNN(dim_in, dim_out)[source]¶

Bases: Module

GNN model that customizes the torch_geometric.graphgym.models.gnn.GNN to support specific handling of new conv layers.

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gnn.layers_pre_mp: Number of pre-message-passing layers. - cfg.gnn.dim_inner: Inner dimension for GNN layers. - cfg.gnn.layers_mp: Number of message-passing layers. - cfg.gnn.dropout: Dropout rate for GNN layers. - cfg.gnn.residual: Whether to use residual connections in GNN layers. - cfg.gnn.layer_type: Type of GNN layer to use (‘gatedgcnconv’ or ‘gineconv’). - cfg.gnn.head: Type of head to use for the final output layer.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.degta_model.DeGTA(dim_in: int, dim_out: int)[source]¶

Bases: Module

Decouplized Graph Triple Attention model. Adapted from https://github.com/wangxiaotang0906/DeGTA

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers: Number of DeGTA layers. - cfg.gt.dim_hidden: Hidden dimension for GNN layers. - cfg.gnn.head: Type of head to use for the final output layer.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.difformer.DIFFormer(dim_in, dim_out)[source]¶

Bases: Module

DIFFormer model. Adapted from https://github.com/qitianwu/DIFFormer

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers: Number of DIFFormer layers. - cfg.gt.dim_hidden: Hidden dimension for GNN layers and DIFFormer layers. - cfg.gnn.head: Type of head to use for the final output layer.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.gps_model.GPSModel(dim_in, dim_out)[source]¶

Bases: Module

General-Powerful-Scalable graph transformer. https://arxiv.org/abs/2205.12454 Rampasek, L., Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., & Beaini, D. Recipe for a general, powerful, scalable graph transformer. (NeurIPS 2022) Adapted from https://github.com/rampasek/GraphGPS

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers (int): Number of GPS layers. - cfg.gt.dim_hidden (int): Hidden dimension for GPS layers. Need to match cfg.gnn.dim_inner. - cfg.gt.layer_type (str): Type of layer to use, containing ‘+’-separated local model type and global model type, e.g., ‘GINE+Transformer’. - cfg.gt.pna_degrees (list): List of PNA degrees for local model. - cfg.gt.n_heads (int): Number of attention heads. - cfg.gt.dropout (float): Dropout rate. - cfg.gt.attn_dropout (float): Attention dropout rate. - cfg.gt.layer_norm (bool): Whether to use layer normalization. - cfg.gt.batch_norm (bool): Whether to use batch normalization. - cfg.gnn.head (str): Type of head to use for the final output layer. - cfg.gnn.layers_pre_mp (int): Number of pre-message-passing layers. - cfg.gnn.dim_inner (int): Inner dimension for GNN layers. Need to match cfg.gt.dim_hidden. - cfg.gnn.act (str): Activation function to use.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.graphormer.GraphormerModel(dim_in, dim_out)[source]¶

Bases: Module

Graphormer port to GraphGPS. https://arxiv.org/abs/2106.05234 Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., … & Liu, T. Y. Do transformers really perform badly for graph representation? (NeurIPS2021) Adapted from https://github.com/rampasek/GraphGPS

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.graphormer.num_layers (int): Number of Graphormer layers. - cfg.graphormer.embed_dim (int): Embedding dimension for Graphormer layers. Need to match cfg.gnn.dim_inner. - cfg.graphormer.num_heads (int): Number of attention heads. - cfg.graphormer.dropout (float): Dropout rate. - cfg.graphormer.attention_dropout (float): Attention dropout rate. - cfg.graphormer.mlp_dropout (float): MLP dropout rate. - cfg.gnn.head (str): Type of head to use for the final output layer. - cfg.gnn.layers_pre_mp (int): Number of pre-message-passing layers. - cfg.gnn.dim_inner (int): Inner dimension for GNN layers. Need to match cfg.graphormer.embed_dim.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.graphtransformer.Graphtransformer(dim_in, dim_out)[source]¶

Bases: Module

Graphtransformer model. Adapted from https://github.com/graphdeeplearning/graphtransformer

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers (int): Number of Graphtransformer layers. - cfg.gt.dim_hidden (int): Hidden dimension for GNN layers and Graphtransformer layers. - cfg.gt.layer_type (str): Type of layer to use for the Graphtransformer layers. - cfg.gnn.head (str): Type of head to use for the final output layer.

Input:

batch (torch_geometric.data.Batch): input batch containing node features and graph structure.

batch.x (torch.Tensor): input node features.
batch.edge_index (torch.Tensor): edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.grit_model.GritTransformer(dim_in, dim_out)[source]¶

Bases: Module

Graph Inductive Bias Transformer (GRIT) model. Adapted from https://github.com/LiamMa/GRIT

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers (int): Number of GRIT layers. - cfg.gt.n_heads (int): Number of attention heads. - cfg.gt.dropout (float): Dropout rate for the GRIT layers. - cfg.gt.dim_hidden (int): Hidden dimension for GNN layers and GRIT layers. Need to match cfg.gnn.dim_inner. - cfg.gt.layer_type (str): Type of layer to use for the GRIT layers. - cfg.gt.attn_dropout (float): Dropout rate for the attention mechanism. - cfg.gt.layer_norm (bool): Whether to use layer normalization. - cfg.gt.batch_norm (bool): Whether to use batch normalization. - cfg.gnn.head (str): Type of head to use for the final output layer. - cfg.gnn.layers_pre_mp (int): Number of pre-message-passing layers. - cfg.gnn.dim_inner (int): Inner dimension for GNN layers. Need to match cfg.gt.dim_hidden.

class opengt.network.mlp_mixer.GraphMLPMixer(dim_in, dim_out)[source]¶

Bases: Module

GraphMLPMixer model. Only supports graph-level tasks. Adapted from https://github.com/XiaoxinHe/Graph-ViT-MLPMixer

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.dim_hidden (int): Hidden dimension for GNN layers and GraphMLPMixer layers. - cfg.gt.mlpmixer_layers (int): Number of Mixer blocks. - cfg.gt.mlpmixer_dropout (float): Dropout rate for the Mixer blocks. - cfg.metis.patches (int): Number of patches. - cfg.metis.patch_rw_dim (int): Dimension of the patch random walk encoder.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.
batch.subgraphs_nodes_mapper (torch.Tensor): Mapping of nodes to subgraphs.
batch.subgraphs_edges_mapper (torch.Tensor): Mapping of edges to subgraphs.
batch.combined_subgraphs (torch.Tensor): Combined subgraphs.
batch.subgraphs_batch (torch.Tensor): Batch indices for subgraphs.
batch.mask (torch.Tensor): Mask for global pooling.

Output:

pred (torch.Tensor): Graph logits after model processing. true (torch.Tensor): Labels for the input batch.

class opengt.network.multi_model.FeatureEncoder(dim_in)[source]¶

Bases: Module

Encoding node and edge features

Parameters:: dim_in (int) – Input feature dimension

class opengt.network.multi_model.MultiModel(dim_in, dim_out)[source]¶

Bases: Module

Multiple layer model for Exphormer and other models. Adapted from https://github.com/hamed1375/Exphormer

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers (int): Number of GPS layers. - cfg.gt.dim_hidden (int): Hidden dimension for GPS layers. Need to match cfg.gnn.dim_inner. - cfg.gt.layer_type (str): Type of layer to use, containing ‘+’-separated local model type and global model type, e.g., ‘GINE+Transformer’. - cfg.gt.pna_degrees (list): List of PNA degrees for local model. - cfg.gt.n_heads (int): Number of attention heads. - cfg.gt.dropout (float): Dropout rate. - cfg.gt.attn_dropout (float): Attention dropout rate. - cfg.gt.layer_norm (bool): Whether to use layer normalization. - cfg.gt.batch_norm (bool): Whether to use batch normalization. - cfg.gnn.head (str): Type of head to use for the final output layer. - cfg.gnn.layers_pre_mp (int): Number of pre-message-passing layers. - cfg.gnn.dim_inner (int): Inner dimension for GNN layers. Need to match cfg.gt.dim_hidden. - cfg.gnn.act (str): Activation function to use.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.multi_model.SingleModel(dim_in, dim_out)[source]¶

Bases: Module

A single layer type can be used without FFN between the layers.

class opengt.network.nodeformer.NodeFormer(dim_in, dim_out)[source]¶

Bases: Module

NodeFormer model. Adapted from https://github.com/qitianwu/NodeFormer Only supports the case where the input is a batch of graphs with the same number of nodes.

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers: Number of NodeFormer layers. - cfg.gt.dim_hidden: Hidden dimension for GNN layers and NodeFormer layers. - cfg.gnn.head: Type of head to use for the final output layer.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.performer.Performer(dim_in, dim_out)[source]¶

Bases: Module

Performer without edge features. This model disregards edge features and runs a linear transformer over a set of node features only. https://arxiv.org/abs/2009.14794 Adapted from https://github.com/rampasek/GraphGPS

class opengt.network.san_transformer.SANTransformer(dim_in, dim_out)[source]¶

Bases: Module

Spectral Attention Network (SAN) Graph Transformer. https://arxiv.org/abs/2106.03893 Adapted from https://github.com/rampasek/GraphGPS

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers (int): Number of SAN layers. - cfg.gt.dim_hidden (int): Hidden dimension for GNN layers and SAN layers. Need to match cfg.gnn.dim_inner. - cfg.gt.gamma (float): Gamma parameter for SAN layers. - cfg.gt.n_heads (int): Number of attention heads. - cfg.gt.full_graph (bool): Whether to use full graph attention. - cfg.gt.dropout (float): Dropout rate for the SAN layers. - cfg.gt.layer_norm (bool): Whether to use layer normalization. - cfg.gt.batch_norm (bool): Whether to use batch normalization. - cfg.gt.residual (bool): Whether to use residual connections. - cfg.gnn.head (str): Type of head to use for the final output layer. - cfg.gnn.layers_pre_mp (int): Number of pre-message-passing layers. - cfg.gnn.dim_inner (int): Inner dimension for GNN layers. Need to match cfg.gt.dim_hidden.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.sgformer.SGFormer(dim_in, dim_out)[source]¶

Bases: Module

SGFormer model. Adapted from https://github.com/qitianwu/SGFormer

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers: Number of SGFormer layers. - cfg.gt.dim_hidden: Hidden dimension for GNN layers and SGFormer layers. - cfg.gt.aggregate: Type of aggregation to use for the graph. e.g., ‘add’ or ‘cat’. - cfg.gt.use_graph: Whether to use graph information. - cfg.gt.graph_weight: Weight for the graph information in the aggregation. - cfg.gnn.head: Type of head to use for the final output layer.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.specformer.SpecFormer(dim_in, dim_out)[source]¶

Bases: Module

SpecFormer model. Adapted from https://github.com/DSL-Lab/Specformer Only supports the case where the input is a batch of graphs with the same number of nodes. Needs preprocessing for LapRaw positional encoding.

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.dim_hidden: Hidden dimension for GNN layers and SpecFormer layers. - cfg.gt.n_heads: Number of attention heads. - cfg.gt.dropout: Dropout rate for the model. - cfg.gt.attn_dropout: Dropout rate for the attention mechanism. - cfg.gnn.head: Type of head to use for the final output layer.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.
batch.EigVals (torch.Tensor): Eigenvalues of the graph Laplacian.
batch.EigVecs (torch.Tensor): Eigenvectors of the graph Laplacian.

Output:

batch (task dependent type, see output head): Output after model processing.

class opengt.network.trans_conv.TransConv(dim_in, dim_out)[source]¶

Bases: Module

TransConv module for SGFormer model. Adapted from https://github.com/qitianwu/SGFormer

Parameters:

dim_in (int) – Number of input features.
dim_out (int) – Number of output features.
cfg (dict) – Configuration dictionary containing model parameters from GraphGym. - cfg.gt.layers: Number of TransConv layers.

Input:

batch (torch_geometric.data.Batch): Input batch containing node features and graph structure.

batch.x (torch.Tensor): Input node features.
batch.edge_index (torch.Tensor): Edge indices of the graph.

Output:

batch (torch_geometric.data.Batch): Output batch after processing.