GINEConv

class dgl.nn.pytorch.conv.GINEConv(apply_func=None, init_eps=0, learn_eps=False)[source]

Bases: Module

Graph Isomorphism Network with Edge Features, introduced by Strategies for Pre-training Graph Neural Networks

\[h_i^{(l+1)} = f_\Theta \left((1 + \epsilon) h_i^{l} + \sum_{j\in\mathcal{N}(i)}\mathrm{ReLU}(h_j^{l} + e_{j,i}^{l})\right)\]

where \(e_{j,i}^{l}\) is the edge feature.

Parameters:
  • apply_func (callable module or None) – The \(f_\Theta\) in the formula. If not None, it will be applied to the updated node features. The default value is None.

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

  • learn_eps (bool, optional) – If True, \(\epsilon\) will be a learnable parameter. Default: False.

Examples

>>> import dgl
>>> import torch
>>> import torch.nn as nn
>>> from dgl.nn import GINEConv
>>> g = dgl.graph(([0, 1, 2], [1, 1, 3]))
>>> in_feats = 10
>>> out_feats = 20
>>> nfeat = torch.randn(g.num_nodes(), in_feats)
>>> efeat = torch.randn(g.num_edges(), in_feats)
>>> conv = GINEConv(nn.Linear(in_feats, out_feats))
>>> res = conv(g, nfeat, efeat)
>>> print(res.shape)
torch.Size([4, 20])
forward(graph, node_feat, edge_feat)[source]

Forward computation.

Parameters:
  • graph (DGLGraph) – The graph.

  • node_feat (torch.Tensor or pair of torch.Tensor) – If a torch.Tensor is given, it is the input feature of shape \((N, D_{in})\) where \(D_{in}\) is size of input feature, \(N\) is the number of nodes. If a pair of torch.Tensor is given, the pair must contain two tensors of shape \((N_{in}, D_{in})\) and \((N_{out}, D_{in})\). If apply_func is not None, \(D_{in}\) should fit the input feature size requirement of apply_func.

  • edge_feat (torch.Tensor) – Edge feature. It is a tensor of shape \((E, D_{in})\) where \(E\) is the number of edges.

Returns:

The output feature of shape \((N, D_{out})\) where \(D_{out}\) is the output feature size of apply_func. If apply_func is None, \(D_{out}\) should be the same as \(D_{in}\).

Return type:

torch.Tensor