Repository: dmlc/dgl

Issue:
EdgeConv message passing
https://github.com/dmlc/dgl/blob/ff64bd0de83f865a8076cf50f69525549c32ca87/python/dgl/nn/pytorch/conv/edgeconv.py#L112

According to the equation, the node h_i itself should be the input of `phi`. Shouldn't this be `edges.dst['x']` instead of `edges.src['x']` ?(I am assuming the message travels from src to dst, and thus dst is the node i which receives messages from neighbours)


Discussion:
Good catch!  I'll fix it.

Resolving patch:
diff --git a/python/dgl/nn/mxnet/conv/edgeconv.py b/python/dgl/nn/mxnet/conv/edgeconv.py
--- a/python/dgl/nn/mxnet/conv/edgeconv.py
+++ b/python/dgl/nn/mxnet/conv/edgeconv.py
@@ -19,12 +19,18 @@ class EdgeConv(nn.Block):
     <https://arxiv.org/pdf/1801.07829>`__".  Can be described as follows:
 
     .. math::
-       h_i^{(l+1)} = \max_{j \in \mathcal{N}(i)} \mathrm{ReLU}(
+       h_i^{(l+1)} = \max_{j \in \mathcal{N}(i)} (
        \Theta \cdot (h_j^{(l)} - h_i^{(l)}) + \Phi \cdot h_i^{(l)})
 
     where :math:`\mathcal{N}(i)` is the neighbor of :math:`i`.
     :math:`\Theta` and :math:`\Phi` are linear layers.
 
+    .. note::
+
+       The original formulation includes a ReLU inside the maximum operator.
+       This is equivalent to first applying a maximum operator then applying
+       the ReLU.
+
     Parameters
     ----------
     in_feat : int
@@ -114,13 +120,6 @@ def __init__(self,
             if batch_norm:
                 self.bn = nn.BatchNorm(in_channels=out_feat)
 
-    def message(self, edges):
-        r"""The message computation function
-        """
-        theta_x = self.theta(edges.dst['x'] - edges.src['x'])
-        phi_x = self.phi(edges.src['x'])
-        return {'e': theta_x + phi_x}
-
     def set_allow_zero_in_degree(self, set_value):
         r"""
 
@@ -182,10 +181,13 @@ def forward(self, g, h):
             h_src, h_dst = expand_as_pair(h, g)
             g.srcdata['x'] = h_src
             g.dstdata['x'] = h_dst
+            g.apply_edges(fn.v_sub_u('x', 'x', 'theta'))
+            g.edata['theta'] = self.theta(g.edata['theta'])
+            g.dstdata['phi'] = self.phi(g.dstdata['x'])
             if not self.batch_norm:
-                g.update_all(self.message, fn.max('e', 'x'))
+                g.update_all(fn.e_add_v('theta', 'phi', 'e'), fn.max('e', 'x'))
             else:
-                g.apply_edges(self.message)
+                g.apply_edges(fn.e_add_v('theta', 'phi', 'e'))
                 g.edata['e'] = self.bn(g.edata['e'])
                 g.update_all(fn.copy_e('e', 'm'), fn.max('m', 'x'))
             return g.dstdata['x']
diff --git a/python/dgl/nn/pytorch/conv/edgeconv.py b/python/dgl/nn/pytorch/conv/edgeconv.py
--- a/python/dgl/nn/pytorch/conv/edgeconv.py
+++ b/python/dgl/nn/pytorch/conv/edgeconv.py
@@ -18,12 +18,18 @@ class EdgeConv(nn.Module):
     <https://arxiv.org/pdf/1801.07829>`__".  Can be described as follows:
 
     .. math::
-       h_i^{(l+1)} = \max_{j \in \mathcal{N}(i)} \mathrm{ReLU}(
+       h_i^{(l+1)} = \max_{j \in \mathcal{N}(i)} (
        \Theta \cdot (h_j^{(l)} - h_i^{(l)}) + \Phi \cdot h_i^{(l)})
 
     where :math:`\mathcal{N}(i)` is the neighbor of :math:`i`.
     :math:`\Theta` and :math:`\Phi` are linear layers.
 
+    .. note::
+
+       The original formulation includes a ReLU inside the maximum operator.
+       This is equivalent to first applying a maximum operator then applying
+       the ReLU.
+
     Parameters
     ----------
     in_feat : int
@@ -105,13 +111,6 @@ def __init__(self,
         if batch_norm:
             self.bn = nn.BatchNorm1d(out_feat)
 
-    def message(self, edges):
-        """The message computation function.
-        """
-        theta_x = self.theta(edges.dst['x'] - edges.src['x'])
-        phi_x = self.phi(edges.src['x'])
-        return {'e': theta_x + phi_x}
-
     def set_allow_zero_in_degree(self, set_value):
         r"""
 
@@ -173,10 +172,13 @@ def forward(self, g, feat):
             h_src, h_dst = expand_as_pair(feat, g)
             g.srcdata['x'] = h_src
             g.dstdata['x'] = h_dst
+            g.apply_edges(fn.v_sub_u('x', 'x', 'theta'))
+            g.edata['theta'] = self.theta(g.edata['theta'])
+            g.dstdata['phi'] = self.phi(g.dstdata['x'])
             if not self.batch_norm:
-                g.update_all(self.message, fn.max('e', 'x'))
+                g.update_all(fn.e_add_v('theta', 'phi', 'e'), fn.max('e', 'x'))
             else:
-                g.apply_edges(self.message)
+                g.apply_edges(fn.e_add_v('theta', 'phi', 'e'))
                 # Although the official implementation includes a per-edge
                 # batch norm within EdgeConv, I choose to replace it with a
                 # global batch norm for a number of reasons:


Test patch: