A2GNNBase(in_dim, hid_dim, num_classes, num_layers=1, adv=False, dropout=0.1, act=F.relu, mode='node', **kwargs)

Bases: Module

Base class for A2GNN.

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.

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

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

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

    Other parameters for the backbone.
domain_classifier(x, alpha)

Domain classification with gradient reversal.

Parameters:
  • x (Tensor) –

    Input features.

  • alpha (float) –

    Gradient reversal scaling parameter.
Returns:
  • Tensor –

    Domain classification logits.
Notes
  • Applies gradient reversal layer
  • Binary domain classification
  • Used only when adv=True
  • Helps in domain adaptation
feat_bottleneck(x, edge_index, batch, prop_nums=30)

Feature extraction through propagation layers.

Parameters:
  • x (Tensor) –

    Node features.

  • edge_index (Tensor) –

    Edge indices.

  • batch (Tensor or None) –

    Batch assignment for graph-level tasks.

  • prop_nums (int, default: 30 ) –

    Number of propagation steps. Default: 30.
Returns:
  • Tensor –

    Processed node/graph features.
Notes

Implementation features:

  • Multiple GCN layers with propagation
  • Activation and dropout after each layer
  • Graph pooling for graph-level tasks
  • Configurable propagation steps
feat_classifier(x, edge_index, batch, prop_nums=1)

Classification layer with propagation.

Parameters:
  • x (Tensor) –

    Node/graph features.

  • edge_index (Tensor) –

    Edge indices.

  • batch (Tensor or None) –

    Batch assignment for graph-level tasks.

  • prop_nums (int, default: 1 ) –

    Number of propagation steps. Default: 1.
Returns:
  • Tensor –

    Classification logits.
Notes
  • Node-level: Uses PropGCNConv with propagation
  • Graph-level: Uses linear classification
  • Single propagation step by default
forward(data, prop_nums)

Forward pass of the A2GNN model.

Parameters:
  • data (Data) –

    Input graph data containing features, edge indices, and batch info.

  • prop_nums (int) –

    Number of propagation steps.
Returns:
  • Tensor –

    Model predictions.
Notes

  • Handles both node and graph-level tasks

  • Two-stage processing:

    • Feature bottleneck with multiple propagations
    • Classification with single propagation