graphnetz.models

models

DGI

DGI(in_channels: int, hidden_channels: int = 512)

Bases: Module

Deep Graph Infomax for unsupervised node representation learning.

References

• 1: Veličković, P. et al. (2019). "Deep Graph Infomax." ICLR.

Source code in src/graphnetz/models/_dgi.py

def __init__(self, in_channels: int, hidden_channels: int = 512) -> None:
    super().__init__()
    self.model = DeepGraphInfomax(
        hidden_channels=hidden_channels,
        encoder=_Encoder(in_channels, hidden_channels),
        summary=lambda z, *_: torch.sigmoid(z.mean(dim=0)),
        corruption=_corruption,
    )

GAT

GAT(
    in_channels: int,
    hidden_channels: int,
    out_channels: int,
    heads: int = 8,
    dropout: float = 0.6,
)

Bases: Module

Two-layer Graph Attention Network.

References

• Velickovic2018: Veličković, P. et al. (2018). "Graph Attention Networks." ICLR.

Source code in src/graphnetz/models/_gat.py

def __init__(
    self,
    in_channels: int,
    hidden_channels: int,
    out_channels: int,
    heads: int = 8,
    dropout: float = 0.6,
) -> None:
    super().__init__()
    self.dropout = dropout
    self.conv1 = GATConv(in_channels, hidden_channels, heads=heads, dropout=dropout)
    self.conv2 = GATConv(hidden_channels * heads, out_channels, heads=1, concat=False, dropout=dropout)

GCN

GCN(in_channels: int, hidden_channels: int, out_channels: int)

Bases: Module

Two-layer Graph Convolutional Network.

Parameters:

Name	Type	Description	Default
in_channels	int	The number of input features.	required
hidden_channels	int	The number of hidden features.	required
out_channels	int	The number of output features.	required

References

• Kipf2017: Kipf, T. N., & Welling, M. (2017). "Semi-Supervised Classification with Graph Convolutional Networks." arXiv:1609.02907.

Source code in src/graphnetz/models/_gcn.py

def __init__(self, in_channels: int, hidden_channels: int, out_channels: int) -> None:
    super().__init__()
    self.conv1 = GCNConv(in_channels, hidden_channels)
    self.conv2 = GCNConv(hidden_channels, out_channels)

GIN

GIN(in_channels: int, hidden_channels: int, out_channels: int, num_layers: int = 3)

Bases: Module

Graph Isomorphism Network for graph-level prediction.

References

• Xu2019: Xu, K. et al. (2019). "How Powerful are Graph Neural Networks?" ICLR.

Source code in src/graphnetz/models/_gin.py

def __init__(self, in_channels: int, hidden_channels: int, out_channels: int, num_layers: int = 3) -> None:
    super().__init__()
    self.convs = torch.nn.ModuleList()
    self.convs.append(GINConv(_mlp(in_channels, hidden_channels), train_eps=True))
    for _ in range(num_layers - 1):
        self.convs.append(GINConv(_mlp(hidden_channels, hidden_channels), train_eps=True))
    self.classifier = Linear(hidden_channels, out_channels)

GraphTransformer

GraphTransformer(
    in_channels: int,
    hidden_channels: int,
    out_channels: int,
    heads: int = 4,
    dropout: float = 0.1,
)

Bases: Module

Two-layer graph transformer based on TransformerConv.

References

• Shi2021: Shi, Y. et al. (2021). "Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification." IJCAI.

Source code in src/graphnetz/models/_graph_transformer.py

def __init__(
    self,
    in_channels: int,
    hidden_channels: int,
    out_channels: int,
    heads: int = 4,
    dropout: float = 0.1,
) -> None:
    super().__init__()
    self.dropout = dropout
    self.conv1 = TransformerConv(in_channels, hidden_channels, heads=heads, dropout=dropout)
    self.conv2 = TransformerConv(hidden_channels * heads, out_channels, heads=1, concat=False, dropout=dropout)

GraphSAGE

GraphSAGE(
    in_channels: int,
    hidden_channels: int,
    out_channels: int,
    aggr: str = "mean",
    dropout: float = 0.5,
)

Bases: Module

Two-layer GraphSAGE for node-level prediction.

References

• Hamilton2017: Hamilton, W. L., Ying, R., & Leskovec, J. (2017). "Inductive Representation Learning on Large Graphs." NeurIPS.

Source code in src/graphnetz/models/_graphsage.py

def __init__(
    self,
    in_channels: int,
    hidden_channels: int,
    out_channels: int,
    aggr: str = "mean",
    dropout: float = 0.5,
) -> None:
    super().__init__()
    self.dropout = dropout
    self.conv1 = SAGEConv(in_channels, hidden_channels, aggr=aggr)
    self.conv2 = SAGEConv(hidden_channels, out_channels, aggr=aggr)

Task type adapters

_adapters

Task-type adapters that turn any node-level encoder into a graph-level classifier/regressor or a Deep Graph Infomax model.

This is the glue that lets GCN, GAT, GraphSAGE, and the Graph Transformer plug into every benchmark task in the library, not just node classification.

GraphLevelWrapper

GraphLevelWrapper(encoder: Module, hidden_channels: int, out_channels: int)

Bases: Module

Wrap a node-level encoder for graph-level prediction.

The encoder is expected to map a PyG Data batch to per-node features of shape [N, hidden_channels]. The wrapper adds a global mean pool over the batch index and a linear classification/regression head.

Source code in src/graphnetz/models/_adapters.py

def __init__(
    self,
    encoder: nn.Module,
    hidden_channels: int,
    out_channels: int,
) -> None:
    super().__init__()
    self.encoder = encoder
    self.head = nn.Linear(hidden_channels, out_channels)

LinkPredWrapper

LinkPredWrapper(encoder: Module)

Bases: Module

Wrap any node-level encoder as a link predictor with a dot-product decoder.

The wrapper exposes encode(data) returning per-node embeddings of shape [N, hidden_channels] and decode(z, edge_label_index) returning a [E] tensor of edge logits.

Source code in src/graphnetz/models/_adapters.py

def __init__(self, encoder: nn.Module) -> None:
    super().__init__()
    self.encoder = encoder

DGIWrapper

DGIWrapper(encoder: Module, hidden_channels: int)

Bases: Module

Wrap any node-level encoder as a Deep Graph Infomax model.

Mirrors the DGI interface (forward(data) returning the (pos_z, neg_z, summary) triple, plus a loss(...) helper) so the benchmark trainer does not need to special-case it.

Source code in src/graphnetz/models/_adapters.py

def __init__(self, encoder: nn.Module, hidden_channels: int) -> None:
    super().__init__()
    self.model = DeepGraphInfomax(
        hidden_channels=hidden_channels,
        encoder=_DataAdapter(encoder),
        summary=lambda z, *_: torch.sigmoid(z.mean(dim=0)),
        corruption=_corruption,
    )