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

Bases: Module

GNN base model.

Parameters:
  • in_dim (int) –

    Input dimension of model.

  • hid_dim (int) –

    Hidden dimension of model.

  • num_classes (int) –

    Number of classes.

  • num_layers (int, default: 1 ) –

    Total number of layers in model. Default: 4.

  • dropout (float, default: 0.1 ) –

    Dropout rate. Default: 0..

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

    Activation function if not None. Default: torch.nn.functional.relu.

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

    The backbone of GNN model. Default: gcn.

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

    Mode for node or graph level tasks. Default: node.

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

    Other parameters for the backbone.
feat_bottleneck(x, edge_index, edge_weight=None)

Feature extraction through GNN layers.

Parameters:
  • x (Tensor) –

    Node feature matrix.

  • edge_index (Tensor) –

    Edge indices.

  • edge_weight (Tensor, default: None ) –

    Edge weights. Default: None.
Returns:
  • Tensor –

    Transformed node features.
Notes

Process:

  1. Sequential GNN layer application
  2. Activation (except last layer)
  3. Dropout regularization
feat_classifier(x, edge_index, edge_weight=None)

Final classification layer.

Parameters:
  • x (Tensor) –

    Node features from bottleneck.

  • edge_index (Tensor) –

    Edge indices.

  • edge_weight (Tensor, default: None ) –

    Edge weights. Default: None.
Returns:
  • Tensor –

    Classification logits.
Notes

Two modes:

  • Node mode: Uses GNN classifier
  • Graph mode: Uses linear classifier
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. Default: None.

  • batch (Tensor, default: None ) –

    Batch vector for graph-level tasks. 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:

  1. Feature transformation through GNN layers
  2. Graph pooling (if graph-level task)
  3. Classification
  4. Log-softmax normalization