libauc.models

This module aims to provide several popular deep neural network implementations collected and adapted from public codebases. We recommend users to cite the original papers when using these models. Here is an overview of this module:

Model

Reference

DenseNet: DenseNet

huang2017densely

ResNet: ResNet

he2016deep

ResNet: ResNet (Cifar version)

he2016deep

NeuMF: NeuMF

he2017neural

MLP: MultiLayer Perceptron

yuan2023libauc (our implementation)

GCN: GCN

kipf2016semi

DeeperGCN: DeeperGCN

li2019deepgcns and li2020deepergcn

GIN: GIN

xu2018powerful

GINE: GINE

hu2019strategies

GAT: GAT

velivckovic2017graph

GAT: GATv2

brody2021attentive

MPNN: MPNN

gilmer2017neural

GraphSAGE: GraphSAGE

hamilton2017inductive

PNA: PNA

corso2020principal

Please refer to the source code for more details about each implementation.

libauc.models.densenet

class DenseNet(growth_rate=32, block_config=(6, 12, 24, 16), num_init_features=64, bn_size=4, drop_rate=0.0, num_classes=1, memory_efficient=False, last_activation=None)[source]

Densenet-BC model class, based on “Densely Connected Convolutional Networks”

Parameters:

  • growth_rate (int) – how many filters to add each layer (k in paper)

  • block_config (list of 4 ints) – how many layers in each pooling block

  • num_init_features (int) – the number of filters to learn in the first convolution layer

  • bn_size (int) – multiplicative factor for number of bottle neck layers (i.e. bn_size * k features in the bottleneck layer)

  • drop_rate (float) – dropout rate after each dense layer

  • num_classes (int) – number of classification classes

  • memory_efficient (bool) – If True, uses checkpointing. Much more memory efficient, but slower. Default: False. See “paper”

densenet121(pretrained=False, progress=True, activations='relu', **kwargs)

Densenet-121 model from “Densely Connected Convolutional Networks”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

  • memory_efficient (bool) –

    but slower. Default: False. See “paper”

densenet161(pretrained=False, progress=True, activations='relu', **kwargs)

Densenet-161 model from “Densely Connected Convolutional Networks”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

  • memory_efficient (bool) –

    but slower. Default: False. See “paper”

densenet169(pretrained=False, progress=True, activations='relu', **kwargs)

Densenet-169 model from “Densely Connected Convolutional Networks”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

  • memory_efficient (bool) –

    but slower. Default: False. See “paper”

densenet201(pretrained=False, progress=True, activations='relu', **kwargs)

Densenet-201 model from “Densely Connected Convolutional Networks”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

  • memory_efficient (bool) –

    but slower. Default: False. See “paper”

libauc.models.neumf

class NeuMF(user_num: int, item_num: int, dropout: float = 0.2, emb_size: int = 64, layers: str = '[64]')[source]

NeuMF is a widely-used model for recommender systems.

Parameters:

  • user_num (int) – the number of users in the dataset

  • item_num (int) – the number of items in the dataset

  • dropout (float, optional) – dropout ratio for the model

  • emb_size (int, optional) – embedding size of the model

  • layers (string, optional) – describe the layer information of the model

Reference:
static init_weights(m)[source]
load_model(model_path=None)[source]
reset_last_layer()[source]
save_model(model_path=None)[source]

libauc.models.perceptron

class MLP(input_dim=29, hidden_sizes=(16,), activation='relu', num_classes=1)[source]

An implementation of Multilayer Perceptron (MLP).

libauc.models.resnet

class ResNet(block, layers, num_classes=1, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None, last_activation=None)[source]
resnet101(pretrained=False, progress=True, activations='relu', **kwargs)[source]

ResNet-101 model from “Deep Residual Learning for Image Recognition”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

resnet152(pretrained=False, progress=True, activations='relu', **kwargs)[source]

ResNet-152 model from “Deep Residual Learning for Image Recognition”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

resnet18(pretrained=False, progress=True, activations='relu', **kwargs)[source]

ResNet-18 model from “Deep Residual Learning for Image Recognition”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

resnet34(pretrained=False, progress=True, activations='relu', **kwargs)[source]

ResNet-34 model from “Deep Residual Learning for Image Recognition”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

resnet50(pretrained=False, progress=True, activations='relu', **kwargs)[source]

ResNet-50 model from “Deep Residual Learning for Image Recognition”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

resnext101_32x8d(pretrained=False, progress=True, activations='relu', **kwargs)[source]

ResNeXt-101 32x8d model from “Aggregated Residual Transformation for Deep Neural Networks”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

resnext50_32x4d(pretrained=False, progress=True, activations='relu', **kwargs)[source]

ResNeXt-50 32x4d model from “Aggregated Residual Transformation for Deep Neural Networks”

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

wide_resnet101_2(pretrained=False, progress=True, activations='relu', **kwargs)[source]

Wide ResNet-101-2 model from “Wide Residual Networks”

The model is the same as ResNet except for the bottleneck number of channels which is twice larger in every block. The number of channels in outer 1x1 convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 channels, and in Wide ResNet-50-2 has 2048-1024-2048.

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

wide_resnet50_2(pretrained=False, progress=True, activations='relu', **kwargs)[source]

Wide ResNet-50-2 model from “Wide Residual Networks”

The model is the same as ResNet except for the bottleneck number of channels which is twice larger in every block. The number of channels in outer 1x1 convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 channels, and in Wide ResNet-50-2 has 2048-1024-2048.

Parameters:

  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

libauc.models.resnet_cifar

class ResNet(block, num_blocks, num_classes=1, last_activation='sigmoid', pretrained=False)[source]
resnet110(pretrained=False, activations='relu', last_activation=None, **kwargs)[source]
resnet1202(pretrained=False, activations='relu', last_activation=None, **kwargs)[source]
resnet20(pretrained=False, activations='relu', last_activation=None, **kwargs)[source]
resnet32(pretrained=False, activations='relu', last_activation=None, **kwargs)[source]
resnet44(pretrained=False, activations='relu', last_activation=None, **kwargs)[source]
resnet56(pretrained=False, activations='relu', last_activation=None, **kwargs)[source]

libauc.models.gnn

class AtomEncoder(emb_dim, atom_features_dims)[source]

Convert discrete atom(node) features to embeddings

class BondEncoder(emb_dim, bond_features_dims)[source]

Convert discrete bond(edge) features to embeddings

class DeeperGCN(num_tasks: int, emb_dim: int, num_layers: int, atom_features_dims: list | None = None, bond_features_dims: list | None = None, graph_pooling='mean', dropout: float = 0.0, aggr: str | List[str] | Aggregation | None = 'softmax', t: float = 0.1, learn_t: bool = True, p: float = 1.0, learn_p: bool = False, block: str = 'res+', act: str | Callable | None = 'relu', norm: str | Callable | None = 'BatchNorm')[source]

The Graph Neural Network from the “DeepGCNs: Can GCNs Go as Deep as CNNs?” paper, using the GENConv operator for message passing. The skip connection operations from the “DeepGCNs: Can GCNs Go as Deep as CNNs?” and “All You Need to Train Deeper GCNs” papers, including the pre-activation residual connection ("res+"), the residual connection ("res"), the dense connection ("dense") and no connections ("plain").

Parameters:

  • num_tasks (int) – Number of tasks for multi-label classification.

  • emb_dim (int) – Size of embedding for each atom(node) and/or bond(edge) and graph.

  • num_layers (int) – Number of message passing layers.

  • atom_features_dims – (list or None, optional): Size of each category feature for atoms(nodes). if a list is not provided, the atom_features_dims will be filled by funtion get_atom_feature_dims().

  • bond_features_dims – (list or None, optional): Size of each category feature for bonds(edges). if a list is not provided, the bond_features_dims will be filled by funtion get_bond_feature_dims().

  • graph_pooling (str) – Pooling function to generate whole-graph embeddings. ("mean", "sum", "max", "attention", "set2set"). (default: "mean")

  • dropout (float, optional) – Dropout probability. (default: 0.)

  • aggr (str or Aggregation, optional) – The aggregation scheme to use. Any aggregation of torch_geometric.nn.aggr can be used, ("softmax", "powermean", "add", "mean", max). (default: "softmax")

  • t (float, optional) – Initial inverse temperature for softmax aggregation. (default: 0.1)

  • learn_t (bool, optional) – If set to True, will learn the value t for softmax aggregation dynamically. (default: True)

  • p (float, optional) – Initial power for power mean aggregation. (default: 1.0)

  • learn_p (bool, optional) – If set to True, will learn the value p for power mean aggregation dynamically. (default: False)

  • block (str, optional) – The skip connection operation to use ("res+", "res", "dense" or "plain"). (default: "res+")

  • act (str or Callable, optional) – The non-linear activation function to use. (default: "relu")

  • norm (str or Callable, optional) – The normalization function to use. (default: "BatchNorm")

class GAT(num_tasks: int, emb_dim: int, num_layers: int, v2=False, atom_features_dims: list | None = None, bond_features_dims: list | None = None, graph_pooling='mean', dropout: float = 0.0, act: str | Callable | None = 'elu', act_first: bool = False, act_kwargs: Dict[str, Any] | None = None, norm: str | Callable | None = 'BatchNorm', norm_kwargs: Dict[str, Any] | None = None, jk: str | None = 'last', **kwargs)[source]

The Graph Neural Network from “Graph Attention Networks” or “How Attentive are Graph Attention Networks?” papers, using the GATConv or GATv2Conv operator for message passing, respectively.

Parameters:

  • num_tasks (int) – Number of tasks for multi-label classification.

  • emb_dim (int) – Size of embedding for each atom(node) and/or bond(edge) and graph.

  • num_layers (int) – Number of message passing layers.

  • v2 (bool, optional) – If set to True, will make use of: class:~torch_geometric.nn.conv.GATv2Conv rather than GATConv. (default: False)

  • atom_features_dims –

    (list or None, optional): Size of each category feature for atoms(nodes). if a list is not provided, the atom_features_dims will be filled by funtion get_atom_feature_dims().

  • bond_features_dims –

    (list or None, optional): Size of each category feature for bonds(edges). if a list is not provided, the bond_features_dims will be filled by funtion get_bond_feature_dims().

  • graph_pooling (str) – Pooling function to generate whole-graph embeddings. ("mean", "sum", "max", "attention", "set2set"). (default: "mean")

  • dropout (float, optional) – Dropout probability. (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: "BatchNorm")

  • norm_kwargs (Dict[str, Any], optional) – Arguments passed to the respective normalization function defined by norm. (default: None)

  • jk (str, optional) – The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (None, "last", "cat", "max", "lstm"). (default: "last")

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.conv.GATConv or torch_geometric.nn.conv.GATv2Conv.

init_conv(in_channels: int | Tuple[int, int], out_channels: int, **kwargs) → MessagePassing[source]
supports_edge_attr: Final[bool] = True
supports_edge_weight: Final[bool] = False
supports_norm_batch: Final[bool]
class GCN(num_tasks: int, emb_dim: int, num_layers: int, atom_features_dims: list | None = None, graph_pooling='mean', dropout: float = 0.5, act: str | Callable | None = 'relu', act_first: bool = False, act_kwargs: Dict[str, Any] | None = None, norm: str | Callable | None = 'BatchNorm', norm_kwargs: Dict[str, Any] | None = None, jk: str | None = 'last', **kwargs)[source]

The Graph Neural Network from the “Semi-supervised Classification with Graph Convolutional Networks” paper, using the GCNConv operator for message passing.

Parameters:

  • num_tasks (int) – Number of tasks for multi-label classification.

  • emb_dim (int) – Size of embedding for each atom and/or bond and graph.

  • num_layers (int) – Number of message passing layers.

  • atom_features_dims –

    (list or None, optional): Size of each category feature for atoms(nodes). if a list is not provided, the atom_features_dims will be filled by funtion get_atom_feature_dims().

  • graph_pooling (str) – Pooling function to generate whole-graph embeddings. ("mean", "sum", "max", "attention", "set2set"). (default: "mean")

  • dropout (float, optional) – Dropout probability. (default: 0.5)

  • 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: "BatchNorm")

  • norm_kwargs (Dict[str, Any], optional) – Arguments passed to the respective normalization function defined by norm. (default: None)

  • jk (str, optional) – The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (None, "last", "cat", "max", "lstm"). (default: "last")

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.conv.GINConv.

init_conv(in_channels: int, out_channels: int, **kwargs) → MessagePassing[source]
supports_edge_attr: Final[bool] = False
supports_edge_weight: Final[bool] = True
supports_norm_batch: Final[bool]
class GIN(num_tasks: int, emb_dim: int, num_layers: int, atom_features_dims: list | None = None, graph_pooling='mean', dropout: float = 0.5, act: str | Callable | None = 'relu', act_first: bool = False, act_kwargs: Dict[str, Any] | None = None, norm: str | Callable | None = 'BatchNorm', norm_kwargs: Dict[str, Any] | None = None, jk: str | None = 'last', **kwargs)[source]

The Graph Neural Network from the “How Powerful are Graph Neural Networks?” paper, using the GINConv operator for message passing.

Parameters:

  • num_tasks (int) – Number of tasks for multi-label classification.

  • emb_dim (int) – Size of embedding for each atom and/or bond and graph.

  • num_layers (int) – Number of message passing layers.

  • atom_features_dims –

    (list or None, optional): Size of each category feature for atoms(nodes). if a list is not provided, the atom_features_dims will be filled by funtion get_atom_feature_dims().

  • graph_pooling (str) – Pooling function to generate whole-graph embeddings. ("mean", "sum", "max", "attention", "set2set"). (default: "mean")

  • dropout (float, optional) – Dropout probability. (default: 0.5)

  • 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: "BatchNorm")

  • norm_kwargs (Dict[str, Any], optional) – Arguments passed to the respective normalization function defined by norm. (default: None)

  • jk (str, optional) – The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (None, "last", "cat", "max", "lstm"). (default: "last")

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.conv.GINConv.

init_conv(in_channels: int, out_channels: int, **kwargs) → MessagePassing[source]
supports_edge_attr: Final[bool] = False
supports_edge_weight: Final[bool] = False
supports_norm_batch: Final[bool]
class GINE(num_tasks: int, emb_dim: int, num_layers: int, atom_features_dims: list | None = None, bond_features_dims: list | None = None, graph_pooling='mean', dropout: float = 0.5, act: str | Callable | None = 'relu', act_first: bool = False, act_kwargs: Dict[str, Any] | None = None, norm: str | Callable | None = 'BatchNorm', norm_kwargs: Dict[str, Any] | None = None, jk: str | None = 'last', **kwargs)[source]

The modified GINConv operator from the “Strategies for Pre-training Graph Neural Networks” paper so that it is able to incorporate edge features \(\mathbf{e}_{j,i}\) into the aggregation procedure.

Parameters:

  • num_tasks (int) – Number of tasks for multi-label classification.

  • emb_dim (int) – Size of embedding for each atom(node) and/or bond(edge) and graph.

  • num_layers (int) – Number of message passing layers.

  • atom_features_dims –

    (list or None, optional): Size of each category feature for atoms(nodes). if a list is not provided, the atom_features_dims will be filled by funtion get_atom_feature_dims().

  • bond_features_dims –

    (list or None, optional): Size of each category feature for bonds(edges). if a list is not provided, the bond_features_dims will be filled by funtion get_bond_feature_dims().

  • graph_pooling (str) – Pooling function to generate whole-graph embeddings. ("mean", "sum", "max", "attention", "set2set"). (default: "mean")

  • dropout (float, optional) – Dropout probability. (default: 0.5)

  • 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: "BatchNorm")

  • norm_kwargs (Dict[str, Any], optional) – Arguments passed to the respective normalization function defined by norm. (default: None)

  • jk (str, optional) – The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (None, "last", "cat", "max", "lstm"). (default: "last")

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.conv.GINEConv.

init_conv(in_channels: int, out_channels: int, **kwargs) → MessagePassing[source]
supports_edge_attr: Final[bool] = True
supports_edge_weight: Final[bool] = False
supports_norm_batch: Final[bool]
class GraphSAGE(num_tasks: int, emb_dim: int, num_layers: int, atom_features_dims: list | None = None, graph_pooling='mean', dropout: float = 0.2, act: str | Callable | None = 'relu', act_first: bool = False, act_kwargs: Dict[str, Any] | None = None, norm: str | Callable | None = 'BatchNorm', norm_kwargs: Dict[str, Any] | None = None, jk: str | None = 'last', **kwargs)[source]

The Graph Neural Network from the “Inductive Representation Learning on Large Graphs” paper, using the SAGEConv operator for message passing.

Parameters:

  • num_tasks (int) – Number of tasks for multi-label classification.

  • emb_dim (int) – Size of embedding for each atom and/or bond and graph.

  • num_layers (int) – Number of message passing layers.

  • atom_features_dims –

    (list or None, optional): Size of each category feature for atoms(nodes). if a list is not provided, the atom_features_dims will be filled by funtion get_atom_feature_dims().

  • graph_pooling (str) – Pooling function to generate whole-graph embeddings. ("mean", "sum", "max", "attention", "set2set"). (default: "mean")

  • dropout (float, optional) – Dropout probability. (default: 0.2)

  • 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: "BatchNorm")

  • norm_kwargs (Dict[str, Any], optional) – Arguments passed to the respective normalization function defined by norm. (default: None)

  • jk (str, optional) – The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (None, "last", "cat", "max", "lstm"). (default: "last")

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.conv.SAGEConv.

init_conv(in_channels: int | Tuple[int, int], out_channels: int, **kwargs) → MessagePassing[source]
supports_edge_attr: Final[bool] = False
supports_edge_weight: Final[bool] = False
supports_norm_batch: Final[bool]
class MPNN(num_tasks: int, emb_dim: int, num_layers: int, atom_features_dims: list | None = None, bond_features_dims: list | None = None, graph_pooling='mean', dropout: float = 0.2, act: str | Callable | None = 'relu', act_first: bool = False, act_kwargs: Dict[str, Any] | None = None, norm: str | Callable | None = 'BatchNorm', norm_kwargs: Dict[str, Any] | None = None, jk: str | None = 'last', **kwargs)[source]

The Graph Neural Network from the “Neural Message Passing for Quantum Chemistry” paper, using the NNConv operator for message passing.

Parameters:

  • num_tasks (int) – Number of tasks for multi-label classification.

  • emb_dim (int) – Size of embedding for each atom(node) and/or bond(edge) and graph.

  • num_layers (int) – Number of message passing layers.

  • atom_features_dims –

    (list or None, optional): Size of each category feature for atoms(nodes). if a list is not provided, the atom_features_dims will be filled by funtion get_atom_feature_dims().

  • bond_features_dims –

    (list or None, optional): Size of each category feature for bonds(edges). if a list is not provided, the bond_features_dims will be filled by funtion get_bond_feature_dims().

  • graph_pooling (str) – Pooling function to generate whole-graph embeddings. ("mean", "sum", "max", "attention", "set2set"). (default: "mean")

  • dropout (float, optional) – Dropout probability. (default: 0.2)

  • 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: "BatchNorm")

  • norm_kwargs (Dict[str, Any], optional) – Arguments passed to the respective normalization function defined by norm. (default: None)

  • jk (str, optional) – The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (None, "last", "cat", "max", "lstm"). (default: "last")

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.conv.NNConv.

init_conv(in_channels: int, out_channels: int, **kwargs) → MessagePassing[source]
supports_edge_attr: Final[bool] = True
supports_edge_weight: Final[bool] = False
supports_norm_batch: Final[bool]
class PNA(num_tasks: int, emb_dim: int, num_layers: int, atom_features_dims: list | None = None, bond_features_dims: list | None = None, graph_pooling='mean', dropout: float = 0.2, act: str | Callable | None = 'relu', act_first: bool = False, act_kwargs: Dict[str, Any] | None = None, norm: str | Callable | None = 'BatchNorm', norm_kwargs: Dict[str, Any] | None = None, jk: str | None = 'last', **kwargs)[source]

The Graph Neural Network from the “Principal Neighbourhood Aggregation for Graph Nets” paper, using the PNAConv operator for message passing.

Parameters:

  • num_tasks (int) – Number of tasks for multi-label classification.

  • emb_dim (int) – Size of embedding for each atom(node) and/or bond(edge) and graph.

  • num_layers (int) – Number of message passing layers.

  • atom_features_dims –

    (list or None, optional): Size of each category feature for atoms(nodes). if a list is not provided, the atom_features_dims will be filled by funtion get_atom_feature_dims().

  • bond_features_dims –

    (list or None, optional): Size of each category feature for bonds(edges). if a list is not provided, the bond_features_dims will be filled by funtion get_bond_feature_dims().

  • graph_pooling (str) – Pooling function to generate whole-graph embeddings. ("mean", "sum", "max", "attention", "set2set"). (default: "mean")

  • dropout (float, optional) – Dropout probability. (default: 0.2)

  • 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: "BatchNorm")

  • norm_kwargs (Dict[str, Any], optional) – Arguments passed to the respective normalization function defined by norm. (default: None)

  • jk (str, optional) – The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (None, "last", "cat", "max", "lstm"). (default: "last")

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.conv.PNAConv.

init_conv(in_channels: int, out_channels: int, **kwargs) → MessagePassing[source]
supports_edge_attr: Final[bool] = True
supports_edge_weight: Final[bool] = False
supports_norm_batch: Final[bool]