GraphATABase(in_dim, hid_dim, num_classes, model_weights, model_list, num_layers=1, dropout=0.1, act=F.relu, gnn='gcn', mode='node', **kwargs)

Bases: Module

Base model for GraphATA.

This is the base implementation of a graph neural network that uses node-centric attention mechanisms for feature transformation and classification. It supports both node-level and graph-level tasks through a flexible architecture.

Parameters:
  • in_dim (int) –

    Input dimension of model (number of node features).

  • hid_dim (int) –

    Hidden dimension of model for intermediate representations.

  • num_classes (int) –

    Number of output classes for classification.

  • model_weights (tuple) –

    Collection of model weights for feature transformation layers.

  • model_list (list) –

    List of models for ensemble classification.

  • num_layers (int, default: 1 ) –

    Total number of GNN layers in model. Default: 1.

  • dropout (float, default: 0.1 ) –

    Dropout rate for regularization. Default: 0.1.

  • act (callable activation function or None, default: relu ) –

    Activation function between GNN layers. Default: torch.nn.functional.relu.

  • gnn (string, default: 'gcn' ) –

    The backbone GNN model type. Default: gcn.

  • mode (str, default: 'node' ) –

    Operation mode, either 'node' for node-level or 'graph' for graph-level tasks. Default: node.

  • **kwargs (optional, default: {} ) –

    Additional parameters for the GNN backbone.
Attributes:
  • in_dim (int) –

    Input feature dimension.

  • hid_dim (int) –

    Hidden layer dimension.

  • num_classes (int) –

    Number of output classes.

  • num_layers (int) –

    Number of GNN layers.

  • dropout (float) –

    Dropout probability.

  • act (callable) –

    Activation function.

  • mode (str) –

    Task mode ('node' or 'graph').

  • convs (ModuleList) –

    List of NodeCentricConv layers.

  • cls (NodeCentricMLP) –

    Classification layer with attention-based ensemble.
Notes

Architecture components:

  • Feature Transformation:

    • Multiple NodeCentricConv layers
    • Activation and dropout between layers
    • Attention-based feature aggregation

  • Task-Specific Processing:

    • Node-level: Direct classification
    • Graph-level: Global pooling before classification

  • Classification:

    • Ensemble of models with attention
    • Task-specific output handling
    • Log-softmax normalization

Examples:

>>> model = GraphATABase(
...     in_dim=64,
...     hid_dim=32,
...     num_classes=7,
...     model_weights=weights,
...     model_list=models,
...     num_layers=2,
...     mode='node'
... )
>>> x = torch.randn(100, 64)  # 100 nodes, 64 features
>>> edge_index = torch.randint(0, 100, (2, 400))  # 400 edges
>>> out = model(x, edge_index)  # Shape: (100, 7)
feat_bottleneck(x, edge_index, edge_weight=None, batch=None)

Feature extraction through GNN layers.

Parameters:
  • x (Tensor) –

    Node feature matrix, shape (num_nodes, in_dim).

  • edge_index (Tensor) –

    Edge indices, shape (2, num_edges).

  • edge_weight (Tensor, default: None ) –

    Edge weights for weighted graph operations. Default: None.

  • batch (Tensor, default: None ) –

    Batch vector for graph-level tasks, indicating node-to-graph assignment. Shape: (num_nodes,). Default: None.
Returns:
  • Tensor –

    Transformed node features, shape (num_nodes, hid_dim).
Notes

Process:

  • Sequential GNN layer application:

    • NodeCentricConv transformation
    • Feature aggregation with attention

  • Regularization:

    • Activation between layers (except last)
    • Dropout for training stability
    • Residual connections through attention
feat_classifier(x, edge_index, edge_weight=None, batch=None)

Final classification layer.

Parameters:
  • x (Tensor) –

    Node features from bottleneck, shape (num_nodes, hid_dim).

  • edge_index (Tensor) –

    Edge indices, shape (2, num_edges).

  • edge_weight (Tensor, default: None ) –

    Edge weights for weighted graph operations. Default: None.

  • batch (Tensor, default: None ) –

    Batch vector for graph-level tasks, indicating node-to-graph assignment. Shape: (num_nodes,). Default: None.
Returns:
  • Tensor –

    Classification logits:

    • Node mode: shape (num_nodes, num_classes)
    • Graph mode: shape (num_graphs, num_classes)
Notes

Operation modes:

  • Node mode:

  • Uses GNN classifier with edge information

  • Maintains graph structure in classification

  • Graph mode:

  • Uses pooled features

  • Global classification without edge information
forward(x, edge_index, edge_weight=None, batch=None)

Forward pass of the GNN model.

Parameters:
  • x (Tensor) –

    Node feature matrix, shape (num_nodes, in_dim).

  • edge_index (Tensor) –

    Edge indices, shape (2, num_edges).

  • edge_weight (Tensor, default: None ) –

    Edge weights for weighted graph operations. Default: None.

  • batch (Tensor, default: None ) –

    Batch vector for graph-level tasks, indicating node-to-graph assignment. Shape: (num_nodes,). Default: None.
Returns:
  • Tensor –

    Log-softmax probabilities:

    • For node mode: shape (num_nodes, num_classes)
    • For graph mode: shape (num_graphs, num_classes)
Notes

Process:

  • Feature transformation through GNN layers:

    • Sequential application of NodeCentricConv
    • Attention-based feature aggregation
    • Activation and dropout between layers

  • Task-specific processing:

    • Node mode: Direct use of node features
    • Graph mode: Global mean pooling over nodes

  • Classification:

    • Attention-based ensemble classification
    • Log-softmax normalization