topobench.nn.readouts.propagate_signal_down module#

Readout layer that propagates the signal from cells of a certain order to the cells of the lower order.

class AbstractZeroCellReadOut(hidden_dim, out_channels, task_level, pooling_type='sum', logits_linear_layer=True, **kwargs)#

Bases: Module

Readout layer for GNNs that operates on the batch level.

Parameters:

hidden_dimint: Hidden dimension of the GNN model.
out_channelsint: Number of output channels.
task_levelstr: Task level for readout layer. Either “graph” or “node”.
pooling_typestr: Pooling type for readout layer. Either “max”, “sum” or “mean”.
logits_linear_layerbool: Whether to use a linear layer for getting the final logits.
**kwargsdict: Additional arguments.

__init__(hidden_dim, out_channels, task_level, pooling_type='sum', logits_linear_layer=True, **kwargs)#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

compute_logits(x, batch)#

Compute logits based on the readout layer.

Parameters:

xtorch.Tensor: Node embeddings.
batchtorch.Tensor: Batch index tensor.

Returns:

torch.Tensor: Logits tensor.

abstract forward(model_out, batch)#

Forward pass.

Parameters:

model_outdict: Dictionary containing the model output.
batchtorch_geometric.data.Data: Batch object containing the batched domain data.

class PropagateSignalDown(**kwargs)#

Bases: AbstractZeroCellReadOut

Propagate signal down readout layer.

This readout layer propagates the signal from cells of a certain order to the cells of the lower order.

Parameters:

**kwargsdict: Additional keyword arguments. It should contain the following keys: - num_cell_dimensions (int): Highest order of cells considered by the model. - self.hidden_dim (int): Dimension of the cells representations. - readout_name (str): Readout name.

__init__(**kwargs)#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(model_out, batch)#

Forward pass of the propagate signal down readout layer.

The layer takes the embeddings of the cells of a certain order and applies a convolutional layer to them. Layer normalization is then applied to the features. The output is concatenated with the initial embeddings of the cells and the result is projected with the use of a linear layer to the dimensions of the cells of lower rank. The process is repeated until the nodes embeddings, which are the cells of rank 0, are reached.

Parameters:

model_outdict: Dictionary containing the model output.
batchtorch_geometric.data.Data: Batch object containing the batched domain data.

Returns:

dict: Dictionary containing the updated model output.