torch_geometric.nn.conv.GINEConv

class GINEConv(nn: Module, eps: float = 0.0, train_eps: bool = False, edge_dim: Optional[int] = None, **kwargs)[source]

Bases: MessagePassing

The modified GINConv operator from the “Strategies for Pre-training Graph Neural Networks” paper.

\[\mathbf{x}^{\prime}_i = h_{\mathbf{\Theta}} \left( (1 + \epsilon) \cdot \mathbf{x}_i + \sum_{j \in \mathcal{N}(i)} \mathrm{ReLU} ( \mathbf{x}_j + \mathbf{e}_{j,i} ) \right)\]

that is able to incorporate edge features \(\mathbf{e}_{j,i}\) into the aggregation procedure.

Parameters:

• nn (torch.nn.Module) – A neural network \(h_{\mathbf{\Theta}}\) that maps node features x of shape [-1, in_channels] to shape [-1, out_channels], e.g., defined by torch.nn.Sequential.

• eps (float, optional) – (Initial) \(\epsilon\)-value. (default: 0.)

• train_eps (bool, optional) – If set to True, \(\epsilon\) will be a trainable parameter. (default: False)

• edge_dim (int, optional) – Edge feature dimensionality. If set to None, node and edge feature dimensionality is expected to match. Other-wise, edge features are linearly transformed to match node feature dimensionality. (default: None)

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

Shapes:

• input: node features \((|\mathcal{V}|, F_{in})\) or \(((|\mathcal{V_s}|, F_{s}), (|\mathcal{V_t}|, F_{t}))\) if bipartite, edge indices \((2, |\mathcal{E}|)\), edge features \((|\mathcal{E}|, D)\) (optional)

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

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

Runs the forward pass of the module.

Return type:

Tensor

reset_parameters()[source]

Resets all learnable parameters of the module.