Repository: pyg-team/pytorch_geometric
Pull request #7644 (2023-06-26 21:37:31)

Clean-up `TGN` example
checked that accuracys match before and after my PR

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

```diff
@@            Coverage Diff             @@
##           master    #7644      +/-   ##
==========================================
- Coverage   91.58%   91.57%   -0.01%     
==========================================
  Files         452      452              
  Lines       25454    25477      +23     
==========================================
+ Hits        23311    23331      +20     
- Misses       2143     2146       +3     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7644?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/loader/temporal\_dataloader.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7644?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2xvYWRlci90ZW1wb3JhbF9kYXRhbG9hZGVyLnB5) | `89.18% <87.50%> (-4.57%)` | :arrow_down: |
| [torch\_geometric/nn/models/tgn.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7644?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy90Z24ucHk=) | `97.09% <100.00%> (+0.03%)` | :arrow_up: |

: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)

puririshi98: @wsad1 @rusty1s thanks for the review, I have addressed your comments and checked that accuracys match before and after my PR. lmk if anything else needed to merge
puririshi98: @rusty1s @wsad1  anything else needed to merge?

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1d4537870..0be489765 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
 
+- Integrate `neg_sampling_ratio` into `TemporalDataLoader` ([#7644](https://github.com/pyg-team/pytorch_geometric/pull/7644))
 - Added `faiss`-based `KNNINdex` classes for L2 or maximum inner product search ([#7842](https://github.com/pyg-team/pytorch_geometric/pull/7842))
 - Added the `OSE_GVCS` dataset ([#7811](https://github.com/pyg-team/pytorch_geometric/pull/7811))
 - Added `output_initializer` argument to `DimeNet` models ([#7774](https://github.com/pyg-team/pytorch_geometric/pull/7774), [#7780](https://github.com/pyg-team/pytorch_geometric/pull/7780))
diff --git a/examples/tgn.py b/examples/tgn.py
index a7c52596d..bd60d40e8 100644
--- a/examples/tgn.py
+++ b/examples/tgn.py
@@ -36,15 +36,24 @@ data = dataset[0]
 # expensive memory transfer costs for mini-batches:
 data = data.to(device)
 
-# Ensure to only sample actual destination nodes as negatives.
-min_dst_idx, max_dst_idx = int(data.dst.min()), int(data.dst.max())
 train_data, val_data, test_data = data.train_val_test_split(
     val_ratio=0.15, test_ratio=0.15)
 
-train_loader = TemporalDataLoader(train_data, batch_size=200)
-val_loader = TemporalDataLoader(val_data, batch_size=200)
-test_loader = TemporalDataLoader(test_data, batch_size=200)
-
+train_loader = TemporalDataLoader(
+    train_data,
+    batch_size=200,
+    neg_sampling_ratio=1.0,
+)
+val_loader = TemporalDataLoader(
+    val_data,
+    batch_size=200,
+    neg_sampling_ratio=1.0,
+)
+test_loader = TemporalDataLoader(
+    test_data,
+    batch_size=200,
+    neg_sampling_ratio=1.0,
+)
 neighbor_loader = LastNeighborLoader(data.num_nodes, size=10, device=device)
 
 
@@ -115,33 +124,25 @@ def train():
 
     total_loss = 0
     for batch in train_loader:
-        batch = batch.to(device)
         optimizer.zero_grad()
+        batch = batch.to(device)
 
-        src, pos_dst, t, msg = batch.src, batch.dst, batch.t, batch.msg
-
-        # Sample negative destination nodes.
-        neg_dst = torch.randint(min_dst_idx, max_dst_idx + 1, (src.size(0), ),
-                                dtype=torch.long, device=device)
-
-        n_id = torch.cat([src, pos_dst, neg_dst]).unique()
-        n_id, edge_index, e_id = neighbor_loader(n_id)
+        n_id, edge_index, e_id = neighbor_loader(batch.n_id)
         assoc[n_id] = torch.arange(n_id.size(0), device=device)
 
         # Get updated memory of all nodes involved in the computation.
         z, last_update = memory(n_id)
         z = gnn(z, last_update, edge_index, data.t[e_id].to(device),
                 data.msg[e_id].to(device))
-
-        pos_out = link_pred(z[assoc[src]], z[assoc[pos_dst]])
-        neg_out = link_pred(z[assoc[src]], z[assoc[neg_dst]])
+        pos_out = link_pred(z[assoc[batch.src]], z[assoc[batch.dst]])
+        neg_out = link_pred(z[assoc[batch.src]], z[assoc[batch.neg_dst]])
 
         loss = criterion(pos_out, torch.ones_like(pos_out))
         loss += criterion(neg_out, torch.zeros_like(neg_out))
 
         # Update memory and neighbor loader with ground-truth state.
-        memory.update_state(src, pos_dst, t, msg)
-        neighbor_loader.insert(src, pos_dst)
+        memory.update_state(batch.src, batch.dst, batch.t, batch.msg)
+        neighbor_loader.insert(batch.src, batch.dst)
 
         loss.backward()
         optimizer.step()
@@ -162,21 +163,15 @@ def test(loader):
     aps, aucs = [], []
     for batch in loader:
         batch = batch.to(device)
-        src, pos_dst, t, msg = batch.src, batch.dst, batch.t, batch.msg
 
-        neg_dst = torch.randint(min_dst_idx, max_dst_idx + 1, (src.size(0), ),
-                                dtype=torch.long, device=device)
-
-        n_id = torch.cat([src, pos_dst, neg_dst]).unique()
-        n_id, edge_index, e_id = neighbor_loader(n_id)
+        n_id, edge_index, e_id = neighbor_loader(batch.n_id)
         assoc[n_id] = torch.arange(n_id.size(0), device=device)
 
         z, last_update = memory(n_id)
         z = gnn(z, last_update, edge_index, data.t[e_id].to(device),
                 data.msg[e_id].to(device))
-
-        pos_out = link_pred(z[assoc[src]], z[assoc[pos_dst]])
-        neg_out = link_pred(z[assoc[src]], z[assoc[neg_dst]])
+        pos_out = link_pred(z[assoc[batch.src]], z[assoc[batch.dst]])
+        neg_out = link_pred(z[assoc[batch.src]], z[assoc[batch.neg_dst]])
 
         y_pred = torch.cat([pos_out, neg_out], dim=0).sigmoid().cpu()
         y_true = torch.cat(
@@ -186,9 +181,8 @@ def test(loader):
         aps.append(average_precision_score(y_true, y_pred))
         aucs.append(roc_auc_score(y_true, y_pred))
 
-        memory.update_state(src, pos_dst, t, msg)
-        neighbor_loader.insert(src, pos_dst)
-
+        memory.update_state(batch.src, batch.dst, batch.t, batch.msg)
+        neighbor_loader.insert(batch.src, batch.dst)
     return float(torch.tensor(aps).mean()), float(torch.tensor(aucs).mean())
 
 
diff --git a/torch_geometric/loader/temporal_dataloader.py b/torch_geometric/loader/temporal_dataloader.py
index d67957935..a67cbffdc 100644
--- a/torch_geometric/loader/temporal_dataloader.py
+++ b/torch_geometric/loader/temporal_dataloader.py
@@ -14,18 +14,31 @@ class TemporalDataLoader(torch.utils.data.DataLoader):
             from which to load the data.
         batch_size (int, optional): How many samples per batch to load.
             (default: :obj:`1`)
+        neg_sampling_ratio (float, optional): The ratio of sampled negative
+            destination nodes to the number of postive destination nodes.
+            (default: :obj:`0.0`)
         **kwargs (optional): Additional arguments of
             :class:`torch.utils.data.DataLoader`.
     """
-    def __init__(self, data: TemporalData, batch_size: int = 1, **kwargs):
+    def __init__(
+        self,
+        data: TemporalData,
+        batch_size: int = 1,
+        neg_sampling_ratio: float = 0.0,
+        **kwargs,
+    ):
         # Remove for PyTorch Lightning:
         kwargs.pop('dataset', None)
         kwargs.pop('collate_fn', None)
-
         kwargs.pop('shuffle', None)
 
         self.data = data
         self.events_per_batch = batch_size
+        self.neg_sampling_ratio = neg_sampling_ratio
+
+        if neg_sampling_ratio > 0:
+            self.min_dst = int(data.dst.min())
+            self.max_dst = int(data.dst.max())
 
         if kwargs.get('drop_last', False) and len(data) % batch_size != 0:
             arange = range(0, len(data) - batch_size, batch_size)
@@ -35,4 +48,20 @@ class TemporalDataLoader(torch.utils.data.DataLoader):
         super().__init__(arange, 1, shuffle=False, collate_fn=self, **kwargs)
 
     def __call__(self, arange: List[int]) -> TemporalData:
-        return self.data[arange[0]:arange[0] + self.events_per_batch]
+        batch = self.data[arange[0]:arange[0] + self.events_per_batch]
+
+        n_ids = [batch.src, batch.dst]
+
+        if self.neg_sampling_ratio > 0:
+            batch.neg_dst = torch.randint(
+                low=self.min_dst,
+                high=self.max_dst + 1,
+                size=(round(self.neg_sampling_ratio * batch.dst.size(0)), ),
+                dtype=batch.dst.dtype,
+                device=batch.dst.device,
+            )
+            n_ids += [batch.neg_dst]
+
+        batch.n_id = torch.cat(n_ids, dim=0).unique()
+
+        return batch


Test patch:
diff --git a/test/nn/models/test_tgn.py b/test/nn/models/test_tgn.py
index cb1f5c6dd..958224599 100644
--- a/test/nn/models/test_tgn.py
+++ b/test/nn/models/test_tgn.py
@@ -1,3 +1,4 @@
+import pytest
 import torch
 
 from torch_geometric.data import TemporalData
@@ -10,7 +11,8 @@ from torch_geometric.nn.models.tgn import (
 )
 
 
-def test_tgn():
+@pytest.mark.parametrize('neg_sampling_ratio', [0.0, 1.0])
+def test_tgn(neg_sampling_ratio):
     memory_dim = 16
     time_dim = 16
 
@@ -20,8 +22,11 @@ def test_tgn():
     msg = torch.randn(10, 16)
     data = TemporalData(src=src, dst=dst, t=t, msg=msg)
 
-    loader = TemporalDataLoader(data, batch_size=5)
-
+    loader = TemporalDataLoader(
+        data,
+        batch_size=5,
+        neg_sampling_ratio=neg_sampling_ratio,
+    )
     neighbor_loader = LastNeighborLoader(data.num_nodes, size=3)
     assert neighbor_loader.cur_e_id == 0
     assert neighbor_loader.e_id.size() == (data.num_nodes, 3)
@@ -39,13 +44,12 @@ def test_tgn():
 
     # Test TGNMemory training:
     for i, batch in enumerate(loader):
-        n_id = torch.cat([batch.src, batch.dst]).unique()
-        n_id, edge_index, e_id = neighbor_loader(n_id)
+        n_id, edge_index, e_id = neighbor_loader(batch.n_id)
         z, last_update = memory(n_id)
         memory.update_state(batch.src, batch.dst, batch.t, batch.msg)
         neighbor_loader.insert(batch.src, batch.dst)
         if i == 0:
-            assert n_id.size(0) == 4
+            assert n_id.size(0) >= 4
             assert edge_index.numel() == 0
             assert e_id.numel() == 0
             assert z.size() == (n_id.size(0), memory_dim)