topobench.nn.encoders.hopse_encoder module#

Encoder class to apply SimpleEncoder.

class AbstractFeatureEncoder#

Bases: Module

Abstract class to define a custom feature encoder.

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

abstractmethod forward(data)#

Forward pass of the feature encoder model.

Parameters:

datatorch_geometric.data.Data: Input data object which should contain x features.

class AtomEncoder(emb_dim)#

Bases: Module

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

forward(x)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class BondEncoder(emb_dim)#

Bases: Module

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

forward(edge_attr)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class HOPSEFeatureEncoder(in_channels, out_channels, proj_dropout=0, selected_dimensions=None, max_hop=3, batch_norm=False, use_atom_encoder=False, use_bond_encoder=False, fuse_pse2cell=False, **kwargs)#

Bases: AbstractFeatureEncoder

Encoder class to apply SimpleEncoder.

The SimpleEncoder is applied to the features of each cell according to a simple

Parameters:

in_channelslist[list[int]]: Input dimensions for the features.
out_channelslist[int]: Output dimensions for the features.
proj_dropoutfloat, optional: Dropout for the BaseEncoders (default: 0).
selected_dimensionslist[int], optional: List of indexes to apply the BaseEncoders to (default: None).
max_hoplist[int], optional: List of indexes to apply the BaseEncoders to in terms of hops (default: None).
batch_normbool, optional: Wether to apply batch normalizaiton when encoding (default: False).
use_atom_encoderbool, optional: If True, replace the encoder for dimension 0 / hop 0 with an OGB AtomEncoder (default: False).
use_bond_encoderbool, optional: If True, replace the encoder for dimension 1 / hop 0 with an OGB BondEncoder (default: False).
fuse_pse2cellbool, optional: If True, concatenate and linearly project per-hop PSE encodings back into the cell features after encoding (default: False).
**kwargsdict, optional: Additional keyword arguments.

__init__(in_channels, out_channels, proj_dropout=0, selected_dimensions=None, max_hop=3, batch_norm=False, use_atom_encoder=False, use_bond_encoder=False, fuse_pse2cell=False, **kwargs)#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(data)#

Forward pass.

The method applies the BaseEncoders to the features of the selected_dimensions.

Parameters:

datatorch_geometric.data.Data: Input data object which should contain x_{i} features for each i in the selected_dimensions.

Returns:

torch_geometric.data.Data: Output data object with updated x_{i} features.

class MLP(channel_list=None, *, in_channels=None, hidden_channels=None, out_channels=None, num_layers=None, dropout=0.0, act='relu', act_first=False, act_kwargs=None, norm='batch_norm', norm_kwargs=None, plain_last=True, bias=True, **kwargs)#

Bases: Module

A Multi-Layer Perception (MLP) model.

There exists two ways to instantiate an MLP:

By specifying explicit channel sizes, e.g.,
```
mlp = MLP([16, 32, 64, 128])
```
creates a three-layer MLP with differently sized hidden layers.

By specifying fixed hidden channel sizes over a number of layers, e.g.,
```
mlp = MLP(in_channels=16, hidden_channels=32,
          out_channels=128, num_layers=3)
```
creates a three-layer MLP with equally sized hidden layers.

Parameters:

channel_list (List[int] or int, optional) – List of input, intermediate and output channels such that len(channel_list) - 1 denotes the number of layers of the MLP (default: None)
in_channels (int, optional) – Size of each input sample. Will override channel_list. (default: None)
hidden_channels (int, optional) – Size of each hidden sample. Will override channel_list. (default: None)
out_channels (int, optional) – Size of each output sample. Will override channel_list. (default: None)
num_layers (int, optional) – The number of layers. Will override channel_list. (default: None)
dropout (float or List[float], optional) – Dropout probability of each hidden embedding. If a list is provided, sets the dropout value per layer. (default: 0.)
act (str or Callable, optional) – The non-linear activation function to use. (default: "relu")
act_first (bool, optional) – If set to True, activation is applied before normalization. (default: False)
act_kwargs (Dict[str, Any], optional) – Arguments passed to the respective activation function defined by act. (default: None)
norm (str or Callable, optional) – The normalization function to use. (default: "batch_norm")
norm_kwargs (Dict[str, Any], optional) – Arguments passed to the respective normalization function defined by norm. (default: None)
plain_last (bool, optional) – If set to False, will apply non-linearity, batch normalization and dropout to the last layer as well. (default: True)
bias (bool or List[bool], optional) – If set to False, the module will not learn additive biases. If a list is provided, sets the bias per layer. (default: True)
**kwargs (optional) – Additional deprecated arguments of the MLP layer.

__init__(channel_list=None, *, in_channels=None, hidden_channels=None, out_channels=None, num_layers=None, dropout=0.0, act='relu', act_first=False, act_kwargs=None, norm='batch_norm', norm_kwargs=None, plain_last=True, bias=True, **kwargs)#

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

forward(x, batch=None, batch_size=None, return_emb=None)#

Forward pass.

Parameters:

x (torch.Tensor) – The source tensor.
batch (torch.Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each element to a specific example. Only needs to be passed in case the underlying normalization layers require the batch information. (default: None)
batch_size (int, optional) – The number of examples \(B\). Automatically calculated if not given. Only needs to be passed in case the underlying normalization layers require the batch information. (default: None)
return_emb (bool, optional) – If set to True, will additionally return the embeddings before execution of the final output layer. (default: False)

reset_parameters()#: Resets all learnable parameters of the module.

property in_channels: int#: Size of each input sample.

property num_layers: int#: The number of layers.

property out_channels: int#: Size of each output sample.

supports_norm_batch: Final[bool]#

class SimpleAtomEncoder(in_channels)#

Bases: Module

Thin wrapper around OGB’s AtomEncoder.

Parameters:

in_channelsint: Embedding dimension passed to the underlying AtomEncoder.

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

forward(x, batch)#

Encode integer atom features.

Parameters:

xtorch.Tensor: Atom feature tensor; will be cast to long before encoding.
batchtorch.Tensor: Batch assignment vector (unused, kept for API compatibility).

Returns:

torch.Tensor: Encoded atom features.

class SimpleBondEncoder(in_channels)#

Bases: Module

Thin wrapper around OGB’s BondEncoder.

Parameters:

in_channelsint: Embedding dimension passed to the underlying BondEncoder.

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

forward(x, batch)#

Encode integer bond features.

Parameters:

xtorch.Tensor: Bond feature tensor; will be cast to long before encoding.
batchtorch.Tensor: Batch assignment vector (unused, kept for API compatibility).

Returns:

torch.Tensor: Encoded bond features.