torch_geometric.nn.models.PNA

class PNA(in_channels: int, hidden_channels: int, num_layers: int, out_channels: Optional[int] = None, dropout: float = 0.0, act: Optional[Union[str, Callable]] = 'relu', act_first: bool = False, act_kwargs: Optional[Dict[str, Any]] = None, norm: Optional[Union[str, Callable]] = None, norm_kwargs: Optional[Dict[str, Any]] = None, jk: Optional[str] = None, **kwargs)[source]

Bases: BasicGNN

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

Parameters:

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

  • hidden_channels (int) – Size of each hidden sample.

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

  • out_channels (int, optional) – If not set to None, will apply a final linear transformation to convert hidden node embeddings to output size out_channels. (default: None)

  • 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: None)

  • 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: None)

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

forward(x: Tensor, edge_index: Union[Tensor, SparseTensor], edge_weight: Optional[Tensor] = None, edge_attr: Optional[Tensor] = None, batch: Optional[Tensor] = None, batch_size: Optional[int] = None, num_sampled_nodes_per_hop: Optional[List[int]] = None, num_sampled_edges_per_hop: Optional[List[int]] = None) → Tensor

Forward pass.

Parameters:

  • x (torch.Tensor) – The input node features.

  • edge_index (torch.Tensor or SparseTensor) – The edge indices.

  • edge_weight (torch.Tensor, optional) – The edge weights (if supported by the underlying GNN layer). (default: None)

  • edge_attr (torch.Tensor, optional) – The edge features (if supported by the underlying GNN layer). (default: None)

  • 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)

  • num_sampled_nodes_per_hop (List[int], optional) – The number of sampled nodes per hop. Useful in NeighborLoader scenarios to only operate on minimal-sized representations. (default: None)

  • num_sampled_edges_per_hop (List[int], optional) – The number of sampled edges per hop. Useful in NeighborLoader scenarios to only operate on minimal-sized representations. (default: None)

reset_parameters()

Resets all learnable parameters of the module.