topobench.nn.backbones package#

Subpackages#

Module contents#

Some models implemented for TopoBenchX with automated exports.

class topobench.nn.backbones.CCCN(in_channels, n_layers=2, dropout=0.0, last_act=False)#

Bases: Module

CCCN model.

Parameters:

in_channelsint: Number of input channels.
n_layersint, optional: Number of layers (default: 2).
dropoutfloat, optional: Dropout rate (default: 0).
last_actbool, optional: If True, the last activation function is applied (default: False).

forward(x, Ld, Lu)#

Forward pass.

Parameters:

xtorch.Tensor: Input tensor.
Ldtorch.Tensor: Domain adjacency matrix.
Lutorch.Tensor: Label adjacency matrix.

Returns:

torch.Tensor: Output tensor.

class topobench.nn.backbones.CW(F_in, F_out)#

Bases: Module

Layer of the CCCN model.

Parameters:

F_inint: Number of input channels.
F_outint: Number of output channels.

forward(xe, Lu, Ld)#

Forward pass.

Parameters:

xetorch.Tensor: Input tensor.
Lutorch.Tensor: Domain adjacency matrix.
Ldtorch.Tensor: Label adjacency matrix.

Returns:

torch.Tensor: Output tensor.

class topobench.nn.backbones.EDGNN(num_features, input_dropout=0.2, dropout=0.2, activation='relu', MLP_num_layers=2, MLP2_num_layers=-1, MLP3_num_layers=-1, All_num_layers=2, edconv_type='EquivSet', restart_alpha=0.5, aggregate='add', normalization='None', AllSet_input_norm=False)#

Bases: Module

EDGNN model.

Parameters:

num_featuresint: Number of input features.
input_dropoutfloat, optional: Dropout rate for input features. Defaults to 0.2.
dropoutfloat, optional: Dropout rate for hidden layers. Defaults to 0.2.
activationstr, optional: Activation function. Defaults to ‘relu’.
MLP_num_layersint, optional: Number of layers in MLP. Defaults to 2.
MLP2_num_layersint, optional: Number of layers in the second MLP. Defaults to -1.
MLP3_num_layersint, optional: Number of layers in the third MLP. Defaults to -1.
All_num_layersint, optional: Number of layers in the EDConv. Defaults to 2.
edconv_typestr, optional: Type of EDConv. Defaults to ‘EquivSet’.
restart_alphafloat, optional: Restart alpha. Defaults to 0.5.
aggregatestr, optional: Aggregation method. Defaults to ‘add’.
normalizationstr, optional: Normalization method. Defaults to ‘None’.
AllSet_input_normbool, optional: Whether to normalize input features. Defaults to False.

forward(x, edge_index)#

Forward pass.

Parameters:

xTensor: Input features.
edge_indexLongTensor: Edge index.

Returns:

Tensor: Output features.
None: None object needed for compatibility.

reset_parameters()#: Reset parameters.

class topobench.nn.backbones.EquivSetConv(in_features, out_features, mlp1_layers=1, mlp2_layers=1, mlp3_layers=1, aggr='add', alpha=0.5, dropout=0.0, normalization='None', input_norm=False)#

Bases: Module

Class implementing the Equivariant Set Convolution.

Parameters:

in_featuresint: Number of input features.
out_featuresint: Number of output features.
mlp1_layersint, optional: Number of layers in the first MLP. Defaults to 1.
mlp2_layersint, optional: Number of layers in the second MLP. Defaults to 1.
mlp3_layersint, optional: Number of layers in the third MLP. Defaults to 1.
aggrstr, optional: Aggregation method. Defaults to ‘add’.
alphafloat, optional: Alpha value. Defaults to 0.5.
dropoutfloat, optional: Dropout rate. Defaults to 0.0.
normalizationstr, optional: Normalization method. Defaults to ‘None’.
input_normbool, optional: Whether to normalize input features. Defaults to False.

forward(X, vertex, edges, X0)#

Forward pass.

Parameters:

XTensor: Input features.
vertexLongTensor: Vertex index.
edgesLongTensor: Edge index.
X0Tensor: Initial features.

Returns:

Tensor: Output features.

reset_parameters()#: Reset parameters.

class topobench.nn.backbones.GraphMLP(in_channels, hidden_channels, order=1, dropout=0.0, **kwargs)#

Bases: Module

“Graph MLP backbone.

Parameters:

in_channelsint: Number of input features.
hidden_channelsint: Number of hidden units.
orderint, optional: To compute order-th power of adj matrix (default: 1).
dropoutfloat, optional: Dropout rate (default: 0.0).
**kwargs: Additional arguments.

forward(x)#

Forward pass.

Parameters:

xtorch.Tensor: Input tensor.

Returns:

torch.Tensor: Output tensor.

class topobench.nn.backbones.IdentityGAT(in_channels, hidden_channels, out_channels, num_layers, norm, heads=1, dropout=0.0)#

Bases: Module

Graph Attention Network (GAT) with identity activation function.

Parameters:

in_channelsint: Number of input features.
hidden_channelsint: Number of hidden units.
out_channelsint: Number of output features.
num_layersint: Number of layers.
normtorch.nn.Module: Normalization layer.
headsint, optional: Number of attention heads. Defaults to 1.
dropoutfloat, optional: Dropout rate. Defaults to 0.0.

forward(x, edge_index)#

Forward pass.

Parameters:

xtorch.Tensor: Input node features.
edge_indextorch.Tensor: Edge indices.

Returns:

torch.Tensor: Output node features.

class topobench.nn.backbones.IdentityGCN(in_channels, hidden_channels, out_channels, num_layers, norm, dropout=0.0)#

Bases: Module

Graph Convolutional Network (GCN) with identity activation function.

Parameters:

in_channelsint: Number of input features.
hidden_channelsint: Number of hidden units.
out_channelsint: Number of output features.
num_layersint: Number of layers.
normtorch.nn.Module: Normalization layer.
dropoutfloat, optional: Dropout rate. Defaults to 0.0.

forward(x, edge_index)#

Forward pass.

Parameters:

xtorch.Tensor: Input node features.
edge_indextorch.Tensor: Edge indices.

Returns:

torch.Tensor: Output node features.

class topobench.nn.backbones.IdentityGIN(in_channels, hidden_channels, out_channels, num_layers, norm, dropout=0.0)#

Bases: Module

Graph Isomorphism Network (GIN) with identity activation function.

Parameters:

in_channelsint: Number of input features.
hidden_channelsint: Number of hidden units.
out_channelsint: Number of output features.
num_layersint: Number of layers.
normtorch.nn.Module: Normalization layer.
dropoutfloat, optional: Dropout rate. Defaults to 0.0.

forward(x, edge_index)#

Forward pass.

Parameters:

xtorch.Tensor: Input node features.
edge_indextorch.Tensor: Edge indices.

Returns:

torch.Tensor: Output node features.

class topobench.nn.backbones.IdentitySAGE(in_channels, hidden_channels, out_channels, num_layers, norm, dropout=0.0)#

Bases: Module

GraphSAGE with identity activation function.

Parameters:

in_channelsint: Number of input features.
hidden_channelsint: Number of hidden units.
out_channelsint: Number of output features.
num_layersint: Number of layers.
normtorch.nn.Module: Normalization layer.
dropoutfloat, optional: Dropout rate. Defaults to 0.0.

forward(x, edge_index)#

Forward pass.

Parameters:

xtorch.Tensor: Input node features.
edge_indextorch.Tensor: Edge indices.

Returns:

torch.Tensor: Output node features.

class topobench.nn.backbones.JumpLinkConv(in_features, out_features, mlp_layers=2, aggr='add', alpha=0.5)#

Bases: Module

Class implementing the JumpLink Convolution.

Parameters:

in_featuresint: Number of input features.
out_featuresint: Number of output features.
mlp_layersint, optional: Number of layers in the MLP. Defaults to 2.
aggrstr, optional: Aggregation method. Defaults to ‘add’.
alphafloat, optional: Alpha value. Defaults to 0.5.

forward(X, vertex, edges, X0, beta=1.0)#

Forward pass.

Parameters:

XTensor: Input features.
vertexLongTensor: Vertex index.
edgesLongTensor: Edge index.
X0Tensor: Initial features.
betafloat, optional: Beta value. Defaults to 1.0.

Returns:

Tensor: Output features.

reset_parameters()#: Reset parameters.

class topobench.nn.backbones.MLP(in_channels, hidden_channels, out_channels, dropout=0.25, norm=None, norm_kwargs=None, act=None, act_kwargs=None, final_act=None, final_act_kwargs=None, num_nodes=None, task_level=None, **kwargs)#

Bases: Module

Multi-Layer Perceptron (MLP).

This class implements a multi-layer perceptron architecture with customizable activation functions and normalization layers.

Parameters:

in_channelsint: The dimensionality of the input features.
hidden_channelsint: The dimensionality of the hidden features.
out_channelsint: The dimensionality of the output features.
dropoutfloat, optional: The dropout rate (default 0.25).
normstr, optional: The normalization layer to use (default None).
norm_kwargsdict, optional: Additional keyword arguments for the normalization layer (default None).
actstr, optional: The activation function to use (default “relu”).
act_kwargsdict, optional: Additional keyword arguments for the activation function (default None).
final_actstr, optional: The final activation function to use (default “sigmoid”).
final_act_kwargsdict, optional: Additional keyword arguments for the final activation function (default None).
num_nodesint, optional: The number of nodes in the input graph (default None).
task_levelint, optional: The task level for the model (default None).
**kwargs: Additional keyword arguments.

build_mlp_layers()#

Build the MLP layers.

Returns:

nn.Sequential: The MLP layers.

build_norm_layers(norm, norm_kwargs)#

Build the normalization layers.

Parameters:

normstr: The normalization layer to use.
norm_kwargsdict: Additional keyword arguments for the normalization layer.

Returns:

list: A list of normalization layers.

forward(x, batch_size)#

Forward pass through the MLP.

Parameters:

xtorch.Tensor: Input tensor.
batch_sizeint: Batch size.

Returns:

torch.Tensor: Output tensor.

class topobench.nn.backbones.MeanDegConv(in_features, out_features, init_features=None, mlp1_layers=1, mlp2_layers=1, mlp3_layers=2)#

Bases: Module

Class implementing the Mean Degree Convolution.

Parameters:

in_featuresint: Number of input features.
out_featuresint: Number of output features.
init_featuresint, optional: Number of initial features. Defaults to None.
mlp1_layersint, optional: Number of layers in the first MLP. Defaults to 1.
mlp2_layersint, optional: Number of layers in the second MLP. Defaults to 1.
mlp3_layersint, optional: Number of layers in the third MLP. Defaults to 2.

forward(X, vertex, edges, X0)#

Forward pass.

Parameters:

XTensor: Input features.
vertexLongTensor: Vertex index.
edgesLongTensor: Edge index.
X0Tensor: Initial features.

Returns:

Tensor: Output features.

reset_parameters()#: Reset parameters.

class topobench.nn.backbones.Mlp(input_dim, hid_dim, dropout)#

Bases: Module

MLP module.

Parameters:

input_dimint: Input dimension.
hid_dimint: Hidden dimension.
dropoutfloat: Dropout rate.

forward(x)#

Forward pass.

Parameters:

xtorch.Tensor: Input tensor.

Returns:

torch.Tensor: Output tensor.

class topobench.nn.backbones.PlainMLP(in_channels, hidden_channels, out_channels, num_layers, dropout=0.5)#

Bases: Module

Class implementing a multi-layer perceptron without normalization.

Adapted from CUAI/CorrectAndSmooth.

Parameters:

in_channelsint: Number of input features.
hidden_channelsint: Number of hidden features.
out_channelsint: Number of output features.
num_layersint: Number of layers.
dropoutfloat, optional: Dropout rate. Defaults to 0.5.

forward(x)#

Forward pass.

Parameters:

xTensor: Input features.

Returns:

Tensor: Output features.

reset_parameters()#: Reset parameters.

class topobench.nn.backbones.SCCNNCustom(in_channels_all, hidden_channels_all, conv_order, sc_order, aggr_norm=False, update_func=None, n_layers=2)#

Bases: Module

SCCNN implementation for complex classification.

Note: In this task, we can consider the output on any order of simplices for the classification task, which of course can be amended by a readout layer.

Parameters:

in_channels_alltuple of int: Dimension of input features on (nodes, edges, faces).
hidden_channels_alltuple of int: Dimension of features of hidden layers on (nodes, edges, faces).
conv_orderint: Order of convolutions, we consider the same order for all convolutions.
sc_orderint: Order of simplicial complex.
aggr_normbool, optional: Whether to normalize the aggregation (default: False).
update_funcstr, optional: Update function for the simplicial complex convolution (default: None).
n_layersint, optional: Number of layers (default: 2).

forward(x_all, laplacian_all, incidence_all)#

Forward computation.

Parameters:

x_alltuple(tensors): Tuple of feature tensors (node, edge, face).
laplacian_alltuple(tensors): Tuple of Laplacian tensors (graph laplacian L0, down edge laplacian L1_d, upper edge laplacian L1_u, face laplacian L2).
incidence_alltuple(tensors): Tuple of order 1 and 2 incidence matrices.

Returns:

tuple(tensors): Tuple of final hidden state tensors (node, edge, face).

class topobench.nn.backbones.SCCNNLayer(in_channels, out_channels, conv_order, sc_order, aggr_norm: bool = False, update_func=None, initialization: str = 'xavier_normal')#

Bases: Module

Layer of a Simplicial Complex Convolutional Neural Network.

Parameters:

in_channelstuple of int: Dimensions of input features on nodes, edges, and faces.
out_channelstuple of int: Dimensions of output features on nodes, edges, and faces.
conv_orderint: Convolution order of the simplicial filters.
sc_orderint: SC order.
aggr_normbool, optional: Whether to normalize the aggregated message by the neighborhood size (default: False).
update_funcstr, optional: Activation function used in aggregation layers (default: None).
initializationstr, optional: Initialization method for the weights (default: “xavier_normal”).

aggr_norm_func(conv_operator, x)#

Perform aggregation normalization.

Parameters:

conv_operatortorch.sparse: Convolution operator.
xtorch.Tensor: Feature tensor.

Returns:

torch.Tensor: Normalized feature tensor.

chebyshev_conv(conv_operator, conv_order, x)#

Perform Chebyshev convolution.

Parameters:

conv_operatortorch.sparse: Convolution operator.
conv_orderint: Order of the convolution.
xtorch.Tensor: Feature tensor.

Returns:

torch.Tensor: Output tensor.

forward(x_all, laplacian_all, incidence_all)#

Forward computation.

Parameters:

x_alltuple of tensors: Tuple of input feature tensors (node, edge, face).
laplacian_alltuple of tensors: Tuple of Laplacian tensors (graph laplacian L0, down edge laplacian L1_d, upper edge laplacian L1_u, face laplacian L2).
incidence_alltuple of tensors: Tuple of order 1 and 2 incidence matrices.

Returns:

torch.Tensor: Output tensor for each 0-cell.
torch.Tensor: Output tensor for each 1-cell.
torch.Tensor: Output tensor for each 2-cell.

reset_parameters(gain: float = 1.414)#

Reset learnable parameters.

Parameters:

gainfloat: Gain for the weight initialization.

update(x)#

Update embeddings on each cell (step 4).

Parameters:

xtorch.Tensor: Input tensor.

Returns:

torch.Tensor: Updated tensor.

class topobench.nn.backbones.TopoTune(GNN, neighborhoods, layers, use_edge_attr, activation)#

Bases: Module

Tunes a GNN model using higher-order relations.

This class takes a GNN and its kwargs as inputs, and tunes it with specified additional relations.

Parameters:

GNNtorch.nn.Module, a class not an object: The GNN class to use. ex: GAT, GCN.
neighborhoodslist of lists: The neighborhoods of interest.
layersint: The number of layers to use. Each layer contains one GNN.
use_edge_attrbool: Whether to use edge attributes.
activationstr: The activation function to use. ex: ‘relu’, ‘tanh’, ‘sigmoid’.

aggregate_inter_nbhd(x_out_per_route)#

Aggregate the outputs of the GNN for each rank.

While the GNN takes care of intra-nbhd aggregation, this will take care of inter-nbhd aggregation. Default: sum.

Parameters:

x_out_per_routedict: The outputs of the GNN for each route.

Returns:

dict: The aggregated outputs of the GNN for each rank.

forward(batch)#

Forward pass of the model.

Parameters:

batchComplex or ComplexBatch(Complex): The input data.

Returns:

dict: The output hidden states of the model per rank.

generate_membership_vectors(batch: Data)#

Generate membership vectors based on batch.cell_statistics.

Parameters:

batchtorch_geometric.data.Data: Batch object containing the batched domain data.

Returns:

dict: The batch membership of the graphs per rank.

get_nbhd_cache(params)#

Cache the nbhd information into a dict for the complex at hand.

Parameters:

paramsdict: The parameters of the batch, containing the complex.

Returns:

dict: The neighborhood cache.

interrank_expand(params, src_rank, dst_rank, nbhd_cache, membership)#

Expand the complex into an interrank Hasse graph.

Parameters:

paramsdict: The parameters of the batch, containting the complex.
src_rankint: The source rank.
dst_rankint: The destination rank.
nbhd_cachedict: The neighborhood cache containing the expanded boundary index and edge attributes.
membershipdict: The batch membership of the graphs per rank.

Returns:

torch_geometric.data.Data: The expanded batch of interrank Hasse graphs for this route.

interrank_gnn_forward(batch_route, layer_idx, route_index, n_dst_cells)#

Forward pass of the GNN (one layer) for an interrank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data: The batch of interrank Hasse graphs for this route.
layer_idxint: The index of the layer.
route_indexint: The index of the route.
n_dst_cellsint: The number of destination cells in the whole batch.

Returns:

torch.tensor: The output of the GNN (updated features).

intrarank_expand(params, src_rank, nbhd)#

Expand the complex into an intrarank Hasse graph.

Parameters:

paramsdict: The parameters of the batch, containting the complex.
src_rankint: The source rank.
nbhdstr: The neighborhood to use.

Returns:

torch_geometric.data.Data: The expanded batch of intrarank Hasse graphs for this route.

intrarank_gnn_forward(batch_route, layer_idx, route_index)#

Forward pass of the GNN (one layer) for an intrarank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data: The batch of intrarank Hasse graphs for this route.
layer_idxint: The index of the TopoTune layer.
route_indexint: The index of the route.

Returns:

torch.tensor: The output of the GNN (updated features).

class topobench.nn.backbones.TopoTune_OneHasse(GNN, neighborhoods, layers, use_edge_attr, activation)#

Bases: Module

Tunes a GNN model using higher-order relations.

This class takes a GNN and its kwargs as inputs, and tunes it with specified additional relations. Unlike the case of TopoTune, this class expects a single Hasse graph as input, where all higher-order neighborhoods are represented as a single adjacency matrix.

Parameters:

GNNtorch.nn.Module, a class not an object: The GNN class to use. ex: GAT, GCN.
neighborhoodslist of lists: The neighborhoods of interest.
layersint: The number of layers to use. Each layer contains one GNN.
use_edge_attrbool: Whether to use edge attributes.
activationstr: The activation function to use. ex: ‘relu’, ‘tanh’, ‘sigmoid’.

aggregate_inter_nbhd(x_out)#

Aggregate the outputs of the GNN for each rank.

While the GNN takes care of intra-nbhd aggregation, this will take care of inter-nbhd aggregation. Default: sum.

Parameters:

x_outtorch.tensor: The output of the GNN, concatenated features of each rank.

Returns:

dict: The aggregated outputs of the GNN for each rank.

all_nbhds_expand(params, membership)#

Expand the complex into a single Hasse graph which contains all ranks and all nbhd.

Parameters:

paramsdict: The parameters of the batch, containing the complex.
membershipdict: The batch membership of the graphs per rank.

Returns:

torch_geometric.data.Data: The expanded Hasse graph.

all_nbhds_gnn_forward(batch_route, layer_idx)#

Forward pass of the GNN (one layer) for an intrarank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data: The batch of intrarank Hasse graphs for this route.
layer_idxint: The index of the TopoTune layer.

Returns:

torch.tensor: The output of the GNN (updated features).

forward(batch)#

Forward pass of the model.

Parameters:

batchComplex or ComplexBatch(Complex): The input data.

Returns:

dict: The output hidden states of the model per rank.

generate_membership_vectors(batch: Data)#

Generate membership vectors based on batch.cell_statistics.

Parameters:

batchtorch_geometric.data.Data: Batch object containing the batched domain data.

Returns:

dict: The batch membership of the graphs per rank.

class topobench.nn.backbones.customMLP(in_channels, hidden_channels, out_channels, num_layers, dropout=0.5, Normalization='bn', InputNorm=False)#

Bases: Module

Class implementing a multi-layer perceptron.

Adapted from CUAI/CorrectAndSmooth

Parameters:

in_channelsint: Number of input features.
hidden_channelsint: Number of hidden features.
out_channelsint: Number of output features.
num_layersint: Number of layers.
dropoutfloat, optional: Dropout rate. Defaults to 0.5.
Normalizationstr, optiona: Normalization method. Defaults to ‘bn’.
InputNormbool, optional: Whether to normalize input features. Defaults to False.

flops(x)#

Calculate FLOPs.

Parameters:

xTensor: Input features.

Returns:

int: FLOPs.

forward(x)#

Forward pass.

Parameters:

xTensor: Input features.

Returns:

Tensor: Output features.

reset_parameters()#: Reset parameters.