Repository: pyg-team/pytorch_geometric
Pull request #7648 (2023-06-27 06:31:15)

Allow usage of `BasicGNN` in `DeepGraphInfomax`
Fixes https://github.com/pyg-team/pytorch_geometric/discussions/7640

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7648?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7648](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7648?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (12a133e) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/fc32568934e75c92fc96e45e3b48b3482812d3f4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (fc32568) will **increase** coverage by `0.02%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7648      +/-   ##
==========================================
+ Coverage   91.58%   91.60%   +0.02%     
==========================================
  Files         452      452              
  Lines       25454    25459       +5     
==========================================
+ Hits        23311    23321      +10     
+ Misses       2143     2138       -5     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7648?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/models/deep\_graph\_infomax.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7648?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy9kZWVwX2dyYXBoX2luZm9tYXgucHk=) | `100.00% <100.00%> (ø)` | |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7648/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)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 97cf602cd..d13eb7fa3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -69,6 +69,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Allowed the usage of `BasicGNN` models in `DeepGraphInfomax` ([#7648](https://github.com/pyg-team/pytorch_geometric/pull/7648))
 - Breaking Change: Made `Data.keys` a method rather than a property ([#7629](https://github.com/pyg-team/pytorch_geometric/pull/7629))
 - Added a `num_edges` parameter to the forward method of `HypergraphConv` ([#7560](https://github.com/pyg-team/pytorch_geometric/pull/7560))
 - Fixed `get_mesh_laplacian` for `normalization="sym"` ([#7544](https://github.com/pyg-team/pytorch_geometric/pull/7544))
diff --git a/examples/infomax_transductive.py b/examples/infomax_transductive.py
index e268d51e9..55670e348 100644
--- a/examples/infomax_transductive.py
+++ b/examples/infomax_transductive.py
@@ -1,7 +1,6 @@
 import os.path as osp
 
 import torch
-import torch.nn as nn
 
 from torch_geometric.datasets import Planetoid
 from torch_geometric.nn import DeepGraphInfomax, GCNConv
@@ -11,11 +10,11 @@ path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset)
 dataset = Planetoid(path, dataset)
 
 
-class Encoder(nn.Module):
+class Encoder(torch.nn.Module):
     def __init__(self, in_channels, hidden_channels):
         super().__init__()
         self.conv = GCNConv(in_channels, hidden_channels, cached=True)
-        self.prelu = nn.PReLU(hidden_channels)
+        self.prelu = torch.nn.PReLU(hidden_channels)
 
     def forward(self, x, edge_index):
         x = self.conv(x, edge_index)
@@ -24,14 +23,16 @@ class Encoder(nn.Module):
 
 
 def corruption(x, edge_index):
-    return x[torch.randperm(x.size(0))], edge_index
+    return x[torch.randperm(x.size(0), device=x.device)], edge_index
 
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = DeepGraphInfomax(
-    hidden_channels=512, encoder=Encoder(dataset.num_features, 512),
-    summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)),
-    corruption=corruption).to(device)
+    hidden_channels=512,
+    encoder=Encoder(dataset.num_features, 512),
+    summary=lambda z, *args, **kwargs: z.mean(dim=0).sigmoid(),
+    corruption=corruption,
+).to(device)
 data = dataset[0].to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
diff --git a/torch_geometric/nn/models/deep_graph_infomax.py b/torch_geometric/nn/models/deep_graph_infomax.py
index ba75f26df..36ff63fdc 100644
--- a/torch_geometric/nn/models/deep_graph_infomax.py
+++ b/torch_geometric/nn/models/deep_graph_infomax.py
@@ -1,3 +1,4 @@
+import copy
 from typing import Callable, Tuple
 
 import torch
@@ -49,10 +50,18 @@ class DeepGraphInfomax(torch.nn.Module):
         """Returns the latent space for the input arguments, their
         corruptions and their summary representation."""
         pos_z = self.encoder(*args, **kwargs)
+
         cor = self.corruption(*args, **kwargs)
         cor = cor if isinstance(cor, tuple) else (cor, )
-        neg_z = self.encoder(*cor)
+        cor_args = cor[:len(args)]
+        cor_kwargs = copy.copy(kwargs)
+        for key, value in zip(kwargs.keys(), cor[len(args):]):
+            cor_kwargs[key] = value
+
+        neg_z = self.encoder(*cor_args, **cor_kwargs)
+
         summary = self.summary(pos_z, *args, **kwargs)
+
         return pos_z, neg_z, summary
 
     def discriminate(self, z: Tensor, summary: Tensor,


Test patch:
diff --git a/test/nn/models/test_deep_graph_infomax.py b/test/nn/models/test_deep_graph_infomax.py
index a3e5fb49f..037c43bb9 100644
--- a/test/nn/models/test_deep_graph_infomax.py
+++ b/test/nn/models/test_deep_graph_infomax.py
@@ -1,34 +1,74 @@
 import torch
 
-from torch_geometric.nn import DeepGraphInfomax
-from torch_geometric.testing import is_full_test
+from torch_geometric.nn import GCN, DeepGraphInfomax
+from torch_geometric.testing import is_full_test, withCUDA
 
 
-def test_deep_graph_infomax():
+@withCUDA
+def test_infomax(device):
     def corruption(z):
         return z + 1
 
-    model = DeepGraphInfomax(hidden_channels=16, encoder=lambda x: x,
-                             summary=lambda z, *args: z.mean(dim=0),
-                             corruption=lambda x: x + 1)
-
+    model = DeepGraphInfomax(
+        hidden_channels=16,
+        encoder=lambda x: x,
+        summary=lambda z, *args: z.mean(dim=0),
+        corruption=lambda x: x + 1,
+    ).to(device)
     assert str(model) == 'DeepGraphInfomax(16)'
 
-    x = torch.ones(20, 16)
+    x = torch.ones(20, 16, device=device)
 
     pos_z, neg_z, summary = model(x)
-    assert pos_z.size() == (20, 16) and neg_z.size() == (20, 16)
+    assert pos_z.size() == (20, 16)
+    assert neg_z.size() == (20, 16)
     assert summary.size() == (16, )
 
+    loss = model.loss(pos_z, neg_z, summary)
+    assert float(loss) >= 0
+
     if is_full_test():
         jit = torch.jit.export(model)
         pos_z, neg_z, summary = jit(x)
         assert pos_z.size() == (20, 16) and neg_z.size() == (20, 16)
         assert summary.size() == (16, )
 
-    loss = model.loss(pos_z, neg_z, summary)
-    assert 0 <= loss.item()
+    acc = model.test(
+        train_z=torch.ones(20, 16),
+        train_y=torch.randint(10, (20, )),
+        test_z=torch.ones(20, 16),
+        test_y=torch.randint(10, (20, )),
+    )
+    assert 0 <= acc <= 1
+
+
+@withCUDA
+def test_infomax_predefined_model(device):
+    def corruption(x, edge_index, edge_weight):
+        return (
+            x[torch.randperm(x.size(0), device=x.device)],
+            edge_index,
+            edge_weight,
+        )
 
-    acc = model.test(torch.ones(20, 16), torch.randint(10, (20, )),
-                     torch.ones(20, 16), torch.randint(10, (20, )))
-    assert 0 <= acc and acc <= 1
+    model = DeepGraphInfomax(
+        hidden_channels=16,
+        encoder=GCN(16, 16, num_layers=2),
+        summary=lambda z, *args, **kwargs: z.mean(dim=0).sigmoid(),
+        corruption=corruption,
+    ).to(device)
+
+    x = torch.randn(4, 16, device=device)
+    edge_index = torch.tensor(
+        [[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]],
+        device=device,
+    )
+    edge_weight = torch.rand(edge_index.size(1), device=device)
+
+    pos_z, neg_z, summary = model(x, edge_index, edge_weight=edge_weight)
+    assert pos_z.size() == (4, 16)
+    assert neg_z.size() == (4, 16)
+    assert summary.size() == (16, )
+
+    loss = model.loss(pos_z, neg_z, summary)
+    assert float(loss) >= 0