rindti.layers

Base layer

class BaseLayer(*args, **kwargs)[source]

Bases: pytorch_lightning.core.lightning.LightningModule

Base class for all layers.

forward(*args, **kwargs)[source]

Node layers

class ChebConvNet(input_dim: int, output_dim: int, hidden_dim: int = 32, K: int = 1, num_layers: int = 4, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Chebyshev Convolution.

Refer to torch_geometric.nn.conv.ChebConv for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden vector. Defaults to 32.

  • K (int, optional) – K parameter. Defaults to 1.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], **kwargs) → torch.Tensor[source]
class FilmConvNet(input_dim: int, output_dim: int, hidden_dim: int = 32, edge_dim: Optional[int] = None, num_layers: int = 10, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

FiLM Convolution.

Refer to torch_geometric.nn.conv.FiLMConv for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden layer(s). Defaults to 32.

  • edge_dim (int, optional) – Size of the edge input vector. Defaults to None.

  • num_layers (int, optional) – Number of layers. Defaults to 10.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], edge_feats: Optional[torch.Tensor] = None, **kwargs) → torch.Tensor[source]
class GatConvNet(input_dim, output_dim: int, hidden_dim: int = 32, heads: int = 4, num_layers: int = 4, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Graph Attention Layer.

Refer to torch_geometric.nn.conv.GATConv for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden vector. Defaults to 32.

  • heads (int, optional) – Number of heads for multi-head attention. Defaults to 4.

  • num_layers (int, optional) – Number of layers. Defaults to 4.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], **kwargs) → torch.Tensor[source]
class GINConvNet(input_dim: int, output_dim: int, hidden_dim: int = 64, num_layers: int = 3, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Graph Isomorphism Network.

Refer to torch_geometric.nn.conv.GINConv for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden vector. Defaults to 32.

  • num_layers (int, optional) – Total number of layers. Defaults to 3.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], **kwargs) → torch.Tensor[source]
class PNAConvNet(input_dim: int, output_dim: int, hidden_dim: int = 32, edge_dim: Optional[int] = None, deg: Optional[torch.Tensor] = None, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Principal Neighborhood Aggregation.

Refer to torch_geometric.nn.conv.PNAConv for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden layer. Defaults to 32.

  • edge_dim (int, optional) – Size of the edge dim. Defaults to None.

  • deg (Tensor, optional) – Degree distribution. Defaults to None.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], **kwargs) → torch.Tensor[source]
class TransformerNet(input_dim, output_dim: int, hidden_dim: int = 32, dropout: float = 0.1, edge_dim: Optional[int] = None, edge_type: str = 'none', heads: int = 1, num_layers: int = 3, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Transformer Network.

Refer to torch_geometric.nn.conv.TransformerConv for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden vector. Defaults to 32.

  • dropout (float, optional) – Dropout probability. Defaults to 0.1.

  • edge_dim (int, optional) – Size of the edge vector. Defaults to None.

  • edge_type (int, optional) – Number of edge types. Defaults to “none.

  • heads (int, optional) – Number of heads. Defaults to 1.

  • num_layers (int, optional) – Number of layers. Defaults to 3.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], edge_feats: Optional[torch.Tensor] = None, **kwargs) → torch.Tensor[source]

Pooling layers

class DiffPoolNet(input_dim: int, output_dim: int, hidden_dim: int = 128, max_nodes: int = 600, dropout: float = 0.2, ratio: float = 0.25, pooling_method: str = 'mincut', **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Differential Pooling module.

Refer to torch_geometric.nn.dense.dense_diff_pool and torch_geometric.nn.dense.dense_mincut_pool for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden vector. Defaults to 32.

  • max_nodes (int, optional) – Maximal number of nodes in a graph. Defaults to 600.

  • dropout (float, optional) – Dropout ratio. Defaults to 0.2.

  • ratio (float, optional) – Pooling ratio. Defaults to 0.25.

  • pooling_method (str, optional) – Type of pooling. Defaults to “mincut”.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], batch: torch.Tensor, **kwargs) → torch.Tensor[source]
class GMTNet(input_dim: int, output_dim: int, hidden_dim: int = 128, ratio: float = 0.25, max_nodes: int = 600, num_heads: int = 4, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Graph Multiset Transformer pooling.

Refer to torch_geometric.nn.glob.GraphMultisetTransformer for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden layer(s). Defaults to 128.

  • ratio (float, optional) – Ratio of the number of nodes to be pooled. Defaults to 0.25.

  • max_nodes (int, optional) – Maximal number of nodes in a graph. Defaults to 600.

  • num_heads (int, optional) – Number of heads. Defaults to 4.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], batch: torch.LongTensor) → torch.Tensor[source]
class MeanPool(*args, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Mean Pooling module. Simply averages the node features.

forward(x: torch.Tensor, edge_index: Union[torch.Tensor, torch_sparse.tensor.SparseTensor], batch: torch.Tensor, **kwargs) → torch.Tensor[source]

Other layers

class MLP(input_dim: int, out_dim: int, hidden_dim: int = 64, num_layers: int = 2, dropout: float = 0.2, **kwargs)[source]

Bases: rindti.layers.base_layer.BaseLayer

Simple Multi-layer perceptron.

Refer to torch.nn.Sequential for more details.

Parameters

  • input_dim (int) – Size of the input vector

  • output_dim (int) – Size of the output vector

  • hidden_dim (int, optional) – Size of the hidden vector. Defaults to 32.

  • num_layers (int, optional) – Total Number of layers. Defaults to 2.

  • dropout (float, optional) – Dropout ratio. Defaults to 0.2.

forward(x: torch.Tensor) → torch.Tensor[source]