torch_geometric.nn.conv.PNAConv

class PNAConv(in_channels: int, out_channels: int, aggregators: List[str], scalers: List[str], deg: Tensor, edge_dim: Optional[int] = None, towers: int = 1, pre_layers: int = 1, post_layers: int = 1, divide_input: bool = False, act: Optional[Union[str, Callable]] = 'relu', act_kwargs: Optional[Dict[str, Any]] = None, train_norm: bool = False, **kwargs)[source]

Bases: MessagePassing

The Principal Neighbourhood Aggregation graph convolution operator from the “Principal Neighbourhood Aggregation for Graph Nets” paper.

\[\mathbf{x}_i^{\prime} = \gamma_{\mathbf{\Theta}} \left( \mathbf{x}_i, \underset{j \in \mathcal{N}(i)}{\bigoplus} h_{\mathbf{\Theta}} \left( \mathbf{x}_i, \mathbf{x}_j \right) \right)\]

with

\[\begin{split}\bigoplus = \underbrace{\begin{bmatrix} 1 \\ S(\mathbf{D}, \alpha=1) \\ S(\mathbf{D}, \alpha=-1) \end{bmatrix} }_{\text{scalers}} \otimes \underbrace{\begin{bmatrix} \mu \\ \sigma \\ \max \\ \min \end{bmatrix}}_{\text{aggregators}},\end{split}\]

where \(\gamma_{\mathbf{\Theta}}\) and \(h_{\mathbf{\Theta}}\) denote MLPs.

Note

For an example of using PNAConv, see examples/pna.py.

Parameters:

  • in_channels (int) – Size of each input sample, or -1 to derive the size from the first input(s) to the forward method.

  • out_channels (int) – Size of each output sample.

  • aggregators (List[str]) – Set of aggregation function identifiers, namely "sum", "mean", "min", "max", "var" and "std".

  • scalers (List[str]) – Set of scaling function identifiers, namely "identity", "amplification", "attenuation", "linear" and "inverse_linear".

  • deg (torch.Tensor) – Histogram of in-degrees of nodes in the training set, used by scalers to normalize.

  • edge_dim (int, optional) – Edge feature dimensionality (in case there are any). (default None)

  • towers (int, optional) – Number of towers (default: 1).

  • pre_layers (int, optional) – Number of transformation layers before aggregation (default: 1).

  • post_layers (int, optional) – Number of transformation layers after aggregation (default: 1).

  • divide_input (bool, optional) – Whether the input features should be split between towers or not (default: False).

  • act (str or callable, optional) – Pre- and post-layer activation function to use. (default: "relu")

  • act_kwargs (Dict[str, Any], optional) – Arguments passed to the respective activation function defined by act. (default: None)

  • train_norm (bool, optional) – Whether normalization parameters are trainable. (default: False)

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

Shapes:

  • input: node features \((|\mathcal{V}|, F_{in})\), edge indices \((2, |\mathcal{E}|)\), edge features \((|\mathcal{E}|, D)\) (optional)

  • output: node features \((|\mathcal{V}|, F_{out})\)

forward(x: Tensor, edge_index: Union[Tensor, SparseTensor], edge_attr: Optional[Tensor] = None) → Tensor[source]

Runs the forward pass of the module.

reset_parameters()[source]

Resets all learnable parameters of the module.

static get_degree_histogram(loader: DataLoader) → Tensor[source]

Returns the degree histogram to be used as input for the deg argument in PNAConv.