Repository: pyg-team/pytorch_geometric
Pull request #7458 (2023-05-29 12:40:09)

Implementation of Directed Graph Neural Networks (Dir-GNN)
Implementation of Directed Graph Neural Networks (Dir-GNN), as introduced in the paper "[Edge Directionality Improves Learning on Heterophilic Graphs](https://arxiv.org/abs/2305.10498)". 

We start by implementing DirSageConv.

We note the following regarding the implementation:

- It's not trivial to implement a generic DirGNN module which just takes a conv as input, and deep copies it to have convolutions in both directions. This is because if we take eg. SageConv as base conv, then we would have the root weights twice, whereas we need them only once. Or for GCN, we actually need to change the normalization to account for both in- and out-degree, and cannot just use two standard GCNConv layers.
- As it is implemented now, DirSageConv uses two SageConv layers with different flows. However, the flow direction "target_to_source" does not support message propagation via torch_sparse.SparseTensor or torch.sparse.Tensor . Should we instead have a generic utility function which is able to transpose any Adj  object, and always use SageConv with source_to_target flow, and feed once the original Adj, and the other time the tranpose?

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7458?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7458](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7458?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (fb9598e) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/d8a651cf0f96c41168593565498734be0a931210?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (d8a651c) will **decrease** coverage by `0.17%`.
> The diff coverage is `96.29%`.

> :exclamation: Current head fb9598e differs from pull request most recent head 3d6d59f. Consider uploading reports for the commit 3d6d59f to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7458      +/-   ##
==========================================
- Coverage   91.60%   91.43%   -0.17%     
==========================================
  Files         445      446       +1     
  Lines       24699    24723      +24     
==========================================
- Hits        22626    22606      -20     
- Misses       2073     2117      +44     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7458?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/conv/dir\_sage\_conv.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7458?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvZGlyX3NhZ2VfY29udi5weQ==) | `96.15% <96.15%> (ø)` | |
| [torch\_geometric/nn/conv/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7458?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |

... and [18 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7458/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

emalgorithm: Thanks @rusty1s, I agree with your point. I implemented a general DirGNN class. In the example I tried it with both GCN and SAGE internal convolution and it works well.

If we generalize the gcn_norm function to handle directed graphs properly, then we should be all set.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f73f18b50..67192843c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `DirGNN` model for learning on directed graphs ([#7458](https://github.com/pyg-team/pytorch_geometric/pull/7458))
 - Allow GPU tensors as input to `NodeLoader` and `LinkLoader` ([#7572](https://github.com/pyg-team/pytorch_geometric/pull/7572))
 - Added an `embedding_device` option to allow for GPU inference in `BasicGNN` ([#7548](https://github.com/pyg-team/pytorch_geometric/pull/7548))
 - Added `Performer` to `GPSConv` and remove `attn_dropout` argument from `GPSConv` ([#7465](https://github.com/pyg-team/pytorch_geometric/pull/7465))
diff --git a/examples/dir_gnn.py b/examples/dir_gnn.py
new file mode 100644
index 000000000..177e819f7
--- /dev/null
+++ b/examples/dir_gnn.py
@@ -0,0 +1,97 @@
+import argparse
+import os.path as osp
+
+import torch
+import torch.nn.functional as F
+
+import torch_geometric.transforms as T
+from torch_geometric.datasets import WikipediaNetwork
+from torch_geometric.nn import DirGNNConv, GCNConv, SAGEConv
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--dataset', type=str, default='chameleon')
+parser.add_argument('--hidden_channels', type=int, default=128)
+parser.add_argument('--lr', type=float, default=0.01)
+parser.add_argument('--epochs', type=int, default=1000)
+parser.add_argument('--alpha', type=float, default=1)
+parser.add_argument('--conv', type=str, default='gcn')
+args = parser.parse_args()
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'Wikipedia')
+dataset = WikipediaNetwork(
+    root=path,
+    name=args.dataset,
+    transform=T.NormalizeFeatures(),
+)
+
+data = dataset[0].to(device)
+data.train_mask = data.train_mask[:, 0]
+data.val_mask = data.val_mask[:, 0]
+data.test_mask = data.test_mask[:, 0]
+
+if args.conv == 'gcn':
+    Conv = GCNConv
+elif args.conv == 'sage':
+    Conv = SAGEConv
+else:
+    raise NotImplementedError
+
+
+class DirGNN(torch.nn.Module):
+    def __init__(self, in_channels, hidden_channels, out_channels, alpha):
+        super().__init__()
+        self.conv1 = Conv(in_channels, hidden_channels)
+        self.conv1 = DirGNNConv(self.conv1, alpha, root_weight=False)
+
+        self.conv2 = Conv(hidden_channels, out_channels)
+        self.conv2 = DirGNNConv(self.conv2, alpha, root_weight=False)
+
+    def forward(self, x, edge_index):
+        x = self.conv1(x, edge_index).relu()
+        x = self.conv2(x, edge_index)
+        return x
+
+
+model = DirGNN(
+    dataset.num_features,
+    args.hidden_channels,
+    dataset.num_classes,
+    alpha=args.alpha,
+).to(device)
+
+optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+
+def train():
+    model.train()
+    optimizer.zero_grad()
+    out = model(data.x, data.edge_index)
+    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
+    loss.backward()
+    optimizer.step()
+    return float(loss)
+
+
+@torch.no_grad()
+def test():
+    model.eval()
+    pred = model(data.x, data.edge_index).argmax(dim=-1)
+
+    accs = []
+    for mask in [data.train_mask, data.val_mask, data.test_mask]:
+        accs.append(int((pred[mask] == data.y[mask]).sum()) / int(mask.sum()))
+    return accs
+
+
+best_val_acc = final_test_acc = 0
+for epoch in range(1, args.epochs + 1):
+    loss = train()
+    train_acc, val_acc, tmp_test_acc = test()
+    if val_acc > best_val_acc:
+        best_val_acc = val_acc
+        test_acc = tmp_test_acc
+
+    print(f'Epoch: {epoch:04d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, '
+          f'Val: {val_acc:.4f}, Test: {test_acc:.4f}')
diff --git a/torch_geometric/nn/conv/__init__.py b/torch_geometric/nn/conv/__init__.py
index 9410e03f0..9a6dae927 100644
--- a/torch_geometric/nn/conv/__init__.py
+++ b/torch_geometric/nn/conv/__init__.py
@@ -59,6 +59,7 @@ from .ssg_conv import SSGConv
 from .point_gnn_conv import PointGNNConv
 from .gps_conv import GPSConv
 from .antisymmetric_conv import AntiSymmetricConv
+from .dir_gnn_conv import DirGNNConv
 
 __all__ = [
     'MessagePassing',
@@ -125,6 +126,7 @@ __all__ = [
     'PointGNNConv',
     'GPSConv',
     'AntiSymmetricConv',
+    'DirGNNConv',
 ]
 
 classes = __all__
diff --git a/torch_geometric/nn/conv/dir_gnn_conv.py b/torch_geometric/nn/conv/dir_gnn_conv.py
new file mode 100644
index 000000000..c14e98d14
--- /dev/null
+++ b/torch_geometric/nn/conv/dir_gnn_conv.py
@@ -0,0 +1,74 @@
+import copy
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn.conv import MessagePassing
+
+
+class DirGNNConv(torch.nn.Module):
+    r"""A generic wrapper for computing graph convolution on directed
+    graphs as described in the `"Edge Directionality Improves Learning on
+    Heterophilic Graphs" <https://arxiv.org/abs/2305.10498>`_ paper.
+    :class:`DirGNNConv` will pass messages both from source nodes to target
+    nodes and from target nodes to source nodes.
+
+    Args:
+        conv (MessagePassing): The underlying
+            :class:`~torch_geometric.nn.conv.MessagePassing` layer to use.
+        alpha (float, optional): The alpha coefficient used to weight the
+            aggregations of in- and out-edges as part of a convex combination.
+            (default: :obj:`0.5`)
+        root_weight (bool, optional): If set to :obj:`True`, the layer will add
+            transformed root node features to the output.
+            (default: :obj:`True`)
+    """
+    def __init__(
+        self,
+        conv: MessagePassing,
+        alpha: float = 0.5,
+        root_weight: bool = True,
+    ):
+        super().__init__()
+
+        self.alpha = alpha
+        self.root_weight = root_weight
+
+        self.conv_in = copy.deepcopy(conv)
+        self.conv_out = copy.deepcopy(conv)
+
+        if hasattr(conv, 'add_self_loops'):
+            self.conv_in.add_self_loops = False
+            self.conv_out.add_self_loops = False
+        if hasattr(conv, 'root_weight'):
+            self.conv_in.root_weight = False
+            self.conv_out.root_weight = False
+
+        if root_weight:
+            self.lin = torch.nn.Linear(conv.in_channels, conv.out_channels)
+        else:
+            self.lin = None
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.conv_in.reset_parameters()
+        self.conv_out.reset_parameters()
+        if self.lin is not None:
+            self.lin.reset_parameters()
+
+    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
+        r""""""
+        x_in = self.conv_in(x, edge_index)
+        x_out = self.conv_out(x, edge_index.flip([0]))
+
+        out = self.alpha * x_out + (1 - self.alpha) * x_in
+
+        if self.root_weight:
+            out = out + self.lin(x)
+
+        return out
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.conv_in}, alpha={self.alpha})'


Test patch:
diff --git a/test/nn/conv/test_dir_gnn_conv.py b/test/nn/conv/test_dir_gnn_conv.py
new file mode 100644
index 000000000..fad36b159
--- /dev/null
+++ b/test/nn/conv/test_dir_gnn_conv.py
@@ -0,0 +1,24 @@
+import torch
+
+from torch_geometric.nn import DirGNNConv, SAGEConv
+
+
+def test_dir_gnn_conv():
+    x = torch.randn(4, 16)
+    edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]])
+
+    conv = DirGNNConv(SAGEConv(16, 32))
+    assert str(conv) == 'DirGNNConv(SAGEConv(16, 32, aggr=mean), alpha=0.5)'
+
+    out = conv(x, edge_index)
+    assert out.size() == (4, 32)
+
+
+def test_static_dir_gnn_conv():
+    x = torch.randn(3, 4, 16)
+    edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]])
+
+    conv = DirGNNConv(SAGEConv(16, 32))
+
+    out = conv(x, edge_index)
+    assert out.size() == (3, 4, 32)