Changes in Temporal Data to support a new Temporal Data Loader

examples/tgn.py

@@ -18,6 +18,7 @@
 from torch.nn import Linear
 
 from torch_geometric.datasets import JODIEDataset
+from torch_geometric.loader import TemporalDataLoader
 from torch_geometric.nn import TGNMemory, TransformerConv
 from torch_geometric.nn.models.tgn import (IdentityMessage, LastAggregator,
                                            LastNeighborLoader)
@@ -26,14 +27,17 @@
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'JODIE')
 dataset = JODIEDataset(path, name='wikipedia')
-data = dataset[0].to(device)
+data = dataset.data
 
 # 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)
+
 neighbor_loader = LastNeighborLoader(data.num_nodes, size=10, device=device)
 
 
@@ -103,7 +107,8 @@ def train():
     neighbor_loader.reset_state()  # Start with an empty graph.
 
     total_loss = 0
-    for batch in train_data.seq_batches(batch_size=200):
+    for batch in train_loader:
+        batch = batch.to(device)
         optimizer.zero_grad()
 
         src, pos_dst, t, msg = batch.src, batch.dst, batch.t, batch.msg
@@ -118,7 +123,8 @@ def train():
 
         # 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], data.msg[e_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]])
@@ -139,15 +145,16 @@ def train():
 
 
 @torch.no_grad()
-def test(inference_data):
+def test(inference_data_loader):
     memory.eval()
     gnn.eval()
     link_pred.eval()
 
     torch.manual_seed(12345)  # Ensure deterministic sampling across epochs.
 
     aps, aucs = [], []
-    for batch in inference_data.seq_batches(batch_size=200):
+    for batch in inference_data_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), ),
@@ -158,7 +165,8 @@ def test(inference_data):
         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], data.msg[e_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]])
@@ -180,7 +188,7 @@ def test(inference_data):
 for epoch in range(1, 51):
     loss = train()
     print(f'  Epoch: {epoch:02d}, Loss: {loss:.4f}')
-    val_ap, val_auc = test(val_data)
-    test_ap, test_auc = test(test_data)
+    val_ap, val_auc = test(val_loader)
+    test_ap, test_auc = test(test_loader)
     print(f' Val AP: {val_ap:.4f},  Val AUC: {val_auc:.4f}')
     print(f'Test AP: {test_ap:.4f}, Test AUC: {test_auc:.4f}')

torch_geometric/data/batch.py

@@ -6,8 +6,8 @@
 import torch
 from torch import Tensor
 
-from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.collate import collate
+from torch_geometric.data.data import BaseData, Data
 from torch_geometric.data.dataset import IndexType
 from torch_geometric.data.separate import separate
 
@@ -54,7 +54,7 @@ class Batch(metaclass=DynamicInheritance):
     :obj:`batch`, which maps each node to its respective graph identifier.
     """
     @classmethod
-    def from_data_list(cls, data_list: Union[List[Data], List[HeteroData]],
+    def from_data_list(cls, data_list: List[BaseData],
                        follow_batch: Optional[List[str]] = None,
                        exclude_keys: Optional[List[str]] = None):
         r"""Constructs a :class:`~torch_geometric.data.Batch` object from a
@@ -80,7 +80,7 @@ def from_data_list(cls, data_list: Union[List[Data], List[HeteroData]],
 
         return batch
 
-    def get_example(self, idx: int) -> Union[Data, HeteroData]:
+    def get_example(self, idx: int) -> BaseData:
         r"""Gets the :class:`~torch_geometric.data.Data` or
         :class:`~torch_geometric.data.HeteroData` object at index :obj:`idx`.
         The :class:`~torch_geometric.data.Batch` object must have been created
@@ -103,8 +103,7 @@ def get_example(self, idx: int) -> Union[Data, HeteroData]:
 
         return data
 
-    def index_select(self,
-                     idx: IndexType) -> Union[List[Data], List[HeteroData]]:
+    def index_select(self, idx: IndexType) -> List[BaseData]:
         r"""Creates a subset of :class:`~torch_geometric.data.Data` or
         :class:`~torch_geometric.data.HeteroData` objects from specified
         indices :obj:`idx`.
@@ -152,7 +151,7 @@ def __getitem__(self, idx: Union[int, np.integer, str, IndexType]) -> Any:
         else:
             return self.index_select(idx)
 
-    def to_data_list(self) -> Union[List[Data], List[HeteroData]]:
+    def to_data_list(self) -> List[BaseData]:
         r"""Reconstructs the list of :class:`~torch_geometric.data.Data` or
         :class:`~torch_geometric.data.HeteroData` objects from the
         :class:`~torch_geometric.data.Batch` object.

torch_geometric/data/temporal.py

@@ -111,17 +111,21 @@ def __prepare_non_str_idx(idx):
                 f'{type(idx).__name__}).')
         return idx
 
+    def __generate_item(self, idx):
+        data = {}
+        num_events = self.num_events
+        for key, item in self._store.items():
+            if item.size(0) == num_events:
+                data[key] = item[idx]
+        return TemporalData(**data)
+
     def __getitem__(self, idx: Any) -> Any:
         if isinstance(idx, str):
             return self._store[idx]
 
         prepared_idx = self.__prepare_non_str_idx(idx)
 
-        data = copy.copy(self)
-        for key, item in data:
-            if item.size(0) == self.num_events:
-                data[key] = item[prepared_idx]
-        return data
+        return self.__generate_item(prepared_idx)
 
     def __setitem__(self, key, value):
         """Sets the attribute :obj:`key` to :obj:`value`."""
@@ -133,7 +137,7 @@ def __delitem__(self, idx):
 
         prepared_idx = self.__prepare_non_str_idx(idx)
 
-        for key, item in self:
+        for key, item in self._store.items():
             if item.shape[0] == self.num_events:
                 del item[prepared_idx]
 
@@ -153,8 +157,11 @@ def __delattr__(self, key: str):
         delattr(self._store, key)
 
     def __iter__(self) -> Iterable:
-        for key, value in self._store.items():
-            yield key, value
+        for idx, _ in enumerate(self.src):
+            yield self.__generate_item(torch.tensor([idx]))
+
+    def __len__(self):
+        return len(self.src)
 
     def __call__(self, *args: List[str]) -> Iterable:
         for key, value in self._store.items(*args):
@@ -231,12 +238,7 @@ def __cat_dim__(self, key: str, value: Any, *args, **kwargs) -> Any:
         return 0
 
     def __inc__(self, key: str, value: Any, *args, **kwargs) -> Any:
-        if 'batch' in key:
-            return int(value.max()) + 1
-        elif key in ['src', 'dst']:
-            return self.num_nodes
-        else:
-            return 0
+        return 0
 
     def __repr__(self) -> str:
         cls = self.__class__.__name__
@@ -247,6 +249,14 @@ def __repr__(self) -> str:
 
     def train_val_test_split(self, val_ratio: float = 0.15,
                              test_ratio: float = 0.15):
+        r"""Split the data in 3 parts: training, validation and test.
+
+        Args:
+            val_ratio (float, optional): The proportion (in percents) of the
+                dataset to include in the validation split. (default: `0.15`)
+            test_ratio (float, optional): The proportion (in percents) of the
+                dataset to include in the test split. (default: `0.15`)
+        """
         val_time, test_time = np.quantile(
             self.t.cpu().numpy(),
             [1. - val_ratio - test_ratio, 1. - test_ratio])
@@ -256,10 +266,6 @@ def train_val_test_split(self, val_ratio: float = 0.15,
 
         return self[:val_idx], self[val_idx:test_idx], self[test_idx:]
 
-    def seq_batches(self, batch_size: int):
-        for start in range(0, self.num_events, batch_size):
-            yield self[start:start + batch_size]
-
     ###########################################################################
 
     def coalesce(self):

torch_geometric/loader/__init__.py

@@ -9,6 +9,7 @@
 from .data_list_loader import DataListLoader
 from .dense_data_loader import DenseDataLoader
 from .neighbor_sampler import NeighborSampler
+from .temporal_dataloader import TemporalDataLoader
 
 __all__ = [
     'DataLoader',
@@ -25,6 +26,7 @@
     'DataListLoader',
     'DenseDataLoader',
     'NeighborSampler',
+    'TemporalDataLoader',
 ]
 
 classes = __all__

torch_geometric/loader/data_list_loader.py

@@ -2,7 +2,8 @@
 
 import torch
 
-from torch_geometric.data import Data, Dataset, HeteroData
+from torch_geometric.data import Dataset
+from torch_geometric.data.data import BaseData
 
 
 def collate_fn(data_list):
@@ -30,7 +31,7 @@ class DataListLoader(torch.utils.data.DataLoader):
             :class:`torch.utils.data.DataLoader`, such as :obj:`drop_last` or
             :obj:`num_workers`.
     """
-    def __init__(self, dataset: Union[Dataset, List[Data], List[HeteroData]],
+    def __init__(self, dataset: Union[Dataset, List[BaseData]],
                  batch_size: int = 1, shuffle: bool = False, **kwargs):
         if 'collate_fn' in kwargs:
             del kwargs['collate_fn']

torch_geometric/loader/dataloader.py

@@ -4,7 +4,8 @@
 import torch.utils.data
 from torch.utils.data.dataloader import default_collate
 
-from torch_geometric.data import Batch, Data, Dataset, HeteroData
+from torch_geometric.data import Batch, Dataset
+from torch_geometric.data.data import BaseData
 
 
 class Collater:
@@ -14,7 +15,7 @@ def __init__(self, follow_batch, exclude_keys):
 
     def __call__(self, batch):
         elem = batch[0]
-        if isinstance(elem, (Data, HeteroData)):
+        if isinstance(elem, BaseData):
             return Batch.from_data_list(batch, self.follow_batch,
                                         self.exclude_keys)
         elif isinstance(elem, torch.Tensor):
@@ -59,7 +60,7 @@ class DataLoader(torch.utils.data.DataLoader):
     """
     def __init__(
         self,
-        dataset: Union[Dataset, List[Data], List[HeteroData]],
+        dataset: Union[Dataset, List[BaseData]],
         batch_size: int = 1,
         shuffle: bool = False,
         follow_batch: Optional[List[str]] = None,

torch_geometric/loader/temporal_dataloader.py

@@ -0,0 +1,43 @@
+from torch.utils.data import _utils
+from torch.utils.data.dataloader import _BaseDataLoaderIter
+
+from torch_geometric.loader import DataLoader
+
+
+class _SingleProcessTemporalDataLoaderIter(_BaseDataLoaderIter):
+    def __init__(self, loader):
+        super(_SingleProcessTemporalDataLoaderIter, self).__init__(loader)
+        assert self._timeout == 0
+        assert self._num_workers == 0
+
+    def _next_data(self):
+        index = self._next_index()  # may raise StopIteration
+        data = self._dataset[index]
+        if self._pin_memory:
+            data = _utils.pin_memory.pin_memory(data)
+        return data
+
+
+class TemporalDataLoader(DataLoader):
+    r"""A data loader which merges data objects from a
+    :class:`torch_geometric.data.TemporalData` to a mini-batch.
+
+    Args:
+        dataset (TemporalData): The :obj:`temporalData` from which to load
+            the data.
+        batch_size (int, optional): How many samples per batch to load.
+            (default: :obj:`1`)
+        shuffle (bool, optional): If set to :obj:`True`, the data will be
+            reshuffled at every epoch. (default: :obj:`False`)
+        follow_batch (List[str], optional): Creates assignment batch
+            vectors for each key in the list. (default: :obj:`None`)
+        exclude_keys (List[str], optional): Will exclude each key in the
+            list. (default: :obj:`None`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch.utils.data.DataLoader`.
+    """
+    def __iter__(self) -> '_BaseDataLoaderIter':
+        if self.num_workers == 0:
+            return _SingleProcessTemporalDataLoaderIter(self)
+        else:
+            return super().__iter__()

torch_geometric/profile/utils.py

@@ -5,13 +5,13 @@
 import subprocess as sp
 import sys
 from collections.abc import Mapping, Sequence
-from typing import Any, Tuple, Union
+from typing import Any, Tuple
 
 import torch
 from torch import Tensor
 from torch_sparse import SparseTensor
 
-from torch_geometric.data import Data, HeteroData
+from torch_geometric.data.data import BaseData
 
 
 def count_parameters(model: torch.nn.Module) -> int:
@@ -36,7 +36,7 @@ def get_model_size(model: torch.nn.Module) -> int:
     return model_size
 
 
-def get_data_size(data: Union[Data, HeteroData]) -> int:
+def get_data_size(data: BaseData) -> int:
     r"""Given a :class:`torch_geometric.data.Data` object, get its theoretical
     memory usage in bytes.