GTrans(in_dim, hid_dim, num_classes, num_layers=3, dropout=0.0, act=F.relu, loop_adj=1, loop_feat=4, ratio=0.1, margin=-1, make_undirected=True, strategy='dropedge', weight_decay=0.0, lr=0.0001, epoch=500, gnn='gcn', device='cuda:0', batch_size=0, num_neigh=-1, verbose=2, **kwargs)

Bases: BaseGDA

Empowering Graph Representation Learning with Test-Time Graph Transformation (ICLR-23).

Parameters:
  • in_dim (int) –

    Input feature dimension.

  • hid_dim (int) –

    Hidden dimension of model.

  • num_classes (int) –

    Total number of classes.

  • num_layers (int, default: 3 ) –

    Total number of layers in model. Default: 3.

  • dropout (float, default: 0.0 ) –

    Dropout rate. Default: 0..

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

    GNN backbone. Default: gcn.

  • loop_adj (int, default: 1 ) –

    Loops for optimizing structure. Default: 1.

  • loop_feat (int, default: 4 ) –

    Loops for optimizing features. Default: 4.

  • ratio (float, default: 0.1 ) –

    Budget B for changing graph structure. Default: 0.1.

  • margin (float, default: -1 ) –

    Test time loss hyperparameter. Default: -1.

  • strategy (str, default: 'dropedge' ) –

    Graph augmentation strategy. Default: dropedge.

  • make_undirected (bool, default: True ) –

    Transform into undirected graph. Default: True.

  • weight_decay (float, default: 0.0 ) –

    Weight decay (L2 penalty). Default: 0..

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

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

  • lr (float, default: 0.0001 ) –

    Learning rate. Default: 0.004.

  • epoch (int, default: 500 ) –

    Maximum number of training epoch. Default: 200.

  • device (str, default: 'cuda:0' ) –

    GPU or CPU. Default: cuda:0.

  • batch_size (int, default: 0 ) –

    Minibatch size, 0 for full batch training. Default: 0.

  • num_neigh (int, default: -1 ) –

    Number of neighbors in sampling, -1 for all neighbors. Default: -1.

  • verbose (int, default: 2 ) –

    Verbosity mode. Range in [0, 3]. Larger value for printing out more log information. Default: 2.

  • **kwargs –

    Other parameters for the model.
augment(feat, edge_index=None, edge_weight=None, p=0.5, strategy='dropedge')

Apply graph augmentation for contrastive learning.

Parameters:
  • feat (Tensor) –

    Node features.

  • edge_index (Tensor, default: None ) –

    Edge indices.

  • edge_weight (Tensor, default: None ) –

    Edge weights.

  • p (float, default: 0.5 ) –

    Augmentation probability.

  • strategy (str, default: 'dropedge' ) –

    Augmentation strategy ('dropedge' or 'shuffle').
Returns:
  • Tensor –

    Augmented graph representations.
Notes

Supports two strategies:

  • dropedge: Randomly drops edges
  • shuffle: Randomly shuffles node features
fit(source_data, target_data)

Train the GTrans model with test-time graph transformation.

Parameters:
  • source_data (Data) –

    Source domain graph data.

  • target_data (Data) –

    Target domain graph data.
Notes

Implementation steps:

Source Training

  • Pretrains model on source domain
  • Freezes model parameters after source training

Target Adaptation

  • Initializes feature and structure perturbations
  • Alternates between feature and structure optimization
  • Uses test-time loss for adaptation
  • Maintains perturbation budget constraints

Edge Sampling

  • Samples final edge structure
  • Ensures connectivity constraints
  • Maintains undirected property if specified
forward_model(data, **kwargs)

Forward pass placeholder for GTrans model.

Parameters:
  • data (Data) –

    Input graph data.

  • **kwargs –

    Additional arguments.
Notes

Placeholder method as GTrans implements custom forward logic through:

  • Source domain training in train_source()
  • Test-time adaptation in fit()
  • Prediction using transformed graphs in predict()
init_model(**kwargs)

Initialize the GTrans base model.

Parameters:
  • **kwargs –

    Additional parameters for model initialization.
Returns:
  • GNNBase –

    Initialized model with specified GNN backbone.
Notes

Configures base GNN model with:

  • Specified backbone architecture (GCN by default)
  • Layer dimensions and dropout
  • Activation functions
predict(data)

Make predictions using transformed graph.

Parameters:
  • data (Data) –

    Input graph data.
Returns:
  • tuple –

    Contains: - logits : torch.Tensor Model predictions. - labels : torch.Tensor True labels.
Notes
  • Uses optimized feature and structure
  • Applies final graph transformation
  • Returns predictions and ground truth
process_graph(data)

Process input graph data.

Parameters:
  • data (Data) –

    Input graph to be processed.
Notes

Placeholder method as graph processing is handled through:

  • Feature perturbation optimization
  • Structure modification via edge sampling
  • Graph augmentation strategies in augment()
  • Test-time transformation procedures
sample_final_edges(n_perturbations, perturbed_edge_weight, data, modified_edge_index, n)

Sample final edge structure based on learned weights.

Parameters:
  • n_perturbations (int) –

    Maximum number of allowed edge modifications.

  • perturbed_edge_weight (Tensor) –

    Learned edge weights.

  • data (Data) –

    Target graph data.

  • modified_edge_index (Tensor) –

    Modified edge indices.

  • n (int) –

    Number of nodes.
Returns:
  • tuple –

    Contains: - edge_index : torch.Tensor Final edge structure. - edge_weight : torch.Tensor Final edge weights.
Notes
  • Uses iterative sampling strategy
  • Maintains best performing structure
  • Ensures perturbation budget constraints
  • Handles undirected graph requirements
test_time_loss(feat, edge_index, edge_weight=None)

Compute test-time adaptation loss.

Parameters:
  • feat (Tensor) –

    Node features.

  • edge_index (Tensor) –

    Edge indices.

  • edge_weight (Tensor, default: None ) –

    Edge weights.
Returns:
  • Tensor –

    Computed loss value.
Notes
  • Implements contrastive learning objective
  • Uses different augmentation strategies
  • Supports margin-based loss variant
  • Compares different graph views
train_source(optimizer)

Train the model on source domain data.

Parameters:
  • optimizer (Optimizer) –

    Optimizer for model parameters.
Notes
  • Performs standard supervised training on source domain
  • Computes cross-entropy loss on labeled data
  • Tracks and logs training metrics
  • Uses full batch training only
update_edge_weights(gradient, optimizer_adj, perturbed_edge_weight)

Update edge weights using gradient information.

Parameters:
  • gradient (Tensor) –

    Computed gradients for edge weights.

  • optimizer_adj (Optimizer) –

    Optimizer for edge weights.

  • perturbed_edge_weight (Tensor) –

    Current edge weights to be updated.
Notes
  • Applies gradient updates to edge weights
  • Maintains minimum weight threshold
  • Uses Adam optimizer for updates