[Example] Temporal Graph Neural Network (#2636)

* Add hgat example

* Add experiment

* Clean code

* clear the code

* Add index in README

* Add index in README

* Add index in README

* Add index in README

* Add index in README

* Add index in README

* Change the code title and folder name

* Ready to merge

* Prepare for rebase and change message passing function

* use git ignore to handle empty file

* change file permission to resolve empty file

* Change permission

* change file mode

* Finish Coding

* working code cpu

* pyg compare

* Accelerate with batching

* FastMode Enabled

* update readme

* Update README.md

* refractor code

* Fix Bug

* add a simple temporal sampling method

* test results

* add train speed result

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* Update README.md

* fix bug

* Fixed Grammar and Format Issue
Co-authored-by: Chen <chesirui@3c22fbe5458c.ant.amazon.com>
Co-authored-by: Tianjun Xiao <xiaotj1990327@gmail.com>
Co-authored-by: Ubuntu <ubuntu@ip-172-31-4-63.ap-northeast-1.compute.internal>
Co-authored-by: WangXuhongCN <wangxuhongcn@163.com>
Co-authored-by: WangXuhongCN <34096632+WangXuhongCN@users.noreply.github.com>

examples/README.md

@@ -81,10 +81,12 @@ The folder contains example implementations of selected research papers related
 | [Dynamic Graph CNN for Learning on Point Clouds](#dgcnnpoint) |  |                                  |                           |                    |                    |
 | [Supervised Community Detection with Line Graph Neural Networks](#lgnn) |  |                                  |                           |                    |                    |
 | [Text Generation from Knowledge Graphs with Graph Transformers](#graphwriter) |  |                                  |                           |                    |                    |
+| [Temporal Graph Networks For Deep Learning on Dynamic Graphs](#tgn) |  | :heavy_check_mark:                                 |                           |                    |                    |
 | [Directional Message Passing for Molecular Graphs](#dimenet) |  |                                  | :heavy_check_mark: |                           |                    |
-|  [Link Prediction Based on Graph Neural Networks](#seal) | | :heavy_check_mark: | | :heavy_check_mark: | :heavy_check_mark: |
+| [Link Prediction Based on Graph Neural Networks](#seal) | | :heavy_check_mark: | | :heavy_check_mark: | :heavy_check_mark: |
 | [Variational Graph Auto-Encoders](#vgae) |  | :heavy_check_mark: | | | |
 | [GNNExplainer: Generating Explanations for Graph Neural Networks](#gnnexplainer) |  :heavy_check_mark: |                                  |                                  |                                  |                                  |
+
 ## 2020
 
 - <a name="mvgrl"></a> Hassani and Khasahmadi. Contrastive Multi-View Representation Learning on Graphs. [Paper link](https://arxiv.org/abs/2006.05582). 
@@ -118,6 +120,10 @@ The folder contains example implementations of selected research papers related
     - Example code: [PyTorch](../examples/pytorch/dimenet)
     - Tags: molecules, molecular property prediction, quantum chemistry
 
+- <a name="dagnn"></a> Rossi et al. Temporal Graph Networks For Deep Learning on Dynamic Graphs. [Paper link](https://arxiv.org/abs/2006.10637).
+    - Example code: [Pytorch](../examples/pytorch/tgn)
+    - Tags: over-smoothing, node classification 
+
 ## 2019
 
 

examples/pytorch/tgn/README.md

+# Temporal Graph Neural Network (TGN)
+
+## DGL Implementation of tgn paper.
+
+This DGL examples implements the GNN mode proposed in the paper [TemporalGraphNeuralNetwork](https://arxiv.org/abs/2006.10637.pdf)
+
+## TGN implementor
+
+This example was implemented by [Ericcsr](https://github.com/Ericcsr) during his SDE internship at the AWS Shanghai AI Lab.
+
+## Graph Dataset
+
+Jodie Wikipedia Temporal dataset. Dataset summary:
+
+- Num Nodes: 9227
+- Num Edges: 157, 474
+- Num Edge Features: 172
+- Edge Feature type: LIWC
+- Time Span: 30 days
+- Chronological Split: Train: 70% Valid: 15% Test: 15%
+
+Jodie Reddit Temporal dataset. Dataset summary:
+
+- Num Nodes: 11,000
+- Num Edges: 672, 447
+- Num Edge Features: 172
+- Edge Feature type: LIWC
+- Time Span: 30 days
+- Chronological Split: Train: 70% Valid: 15% Test: 15%
+
+## How to run example files
+
+In tgn folder, run
+
+**please use `train.py`**
+
+```python
+python train.py --dataset wikipedia
+```
+
+If you want to run in fast mode:
+
+```python
+python train.py --dataset wikipedia --fast_mode
+```
+
+If you want to run in simple mode:
+
+```python
+python train.py --dataset wikipedia --simple_mode
+```
+
+If you want to change memory updating module:
+
+```python
+python train.py --dataset wikipedia --memory_updater [rnn/gru]
+```
+
+## Performance
+
+#### Without New Node in test set
+
+| Models/Datasets | Wikipedia          | Reddit           |
+| --------------- | ------------------ | ---------------- |
+| TGN simple mode | AP: 98.5 AUC: 98.9 | AP: N/A AUC: N/A |
+| TGN fast mode   | AP: 98.2 AUC: 98.6 | AP: N/A AUC: N/A |
+| TGN             | AP: 98.9 AUC: 98.5 | AP: N/A AUC: N/A |
+
+#### With New Node in test set
+
+| Models/Datasets | Wikipedia           | Reddit           |
+| --------------- | ------------------- | ---------------- |
+| TGN simple mode | AP: 98.2  AUC: 98.6 | AP: N/A AUC: N/A |
+| TGN fast mode   | AP: 98.0  AUC: 98.4 | AP: N/A AUC: N/A |
+| TGN             | AP: 98.2  AUC: 98.1 | AP: N/A AUC: N/A |
+
+## Training Speed / Batch
+Intel E5 2cores, Tesla K80, Wikipedia Dataset
+
+| Models/Datasets | Wikipedia | Reddit   |
+| --------------- | --------- | -------- |
+| TGN simple mode | 0.3s      | N/A      |
+| TGN fast mode   | 0.28s     | N/A      |
+| TGN             | 1.3s      | N/A      |
+
+### Details explained
+
+**What is Simple Mode**
+
+Simple Temporal Sampler just choose the edges that happen before the current timestamp and build the subgraph of the corresponding nodes. 
+And then the simple sampler uses the static graph neighborhood sampling methods.
+
+**What is Fast Mode**
+
+Normally temporal encoding needs each node to use incoming time frame as current time which might lead to two nodes have multiple interactions within the same batch need to maintain multiple embedding features which slow down the batching process to avoid feature duplication, fast mode enables fast batching since it uses last memory update time in the last batch as temporal encoding benchmark for each node. Also within each batch, all interaction between two nodes are predicted using the same set of embedding feature
+
+**What is New Node test**
+
+To test the model has the ability to predict link between unseen nodes based on neighboring information of seen nodes. This model deliberately select 10 % of node in test graph and mask them out during the training
+

examples/pytorch/tgn/data_preprocess.py

+import os
+import ssl
+from six.moves import urllib
+
+import pandas as pd
+import numpy as np
+
+import torch
+import dgl
+
+# === Below data preprocessing code are based on
+# https://github.com/twitter-research/tgn
+
+# Preprocess the raw data split each features
+
+def preprocess(data_name):
+    u_list, i_list, ts_list, label_list = [], [], [], []
+    feat_l = []
+    idx_list = []
+
+    with open(data_name) as f:
+        s = next(f)
+        for idx, line in enumerate(f):
+            e = line.strip().split(',')
+            u = int(e[0])
+            i = int(e[1])
+
+            ts = float(e[2])
+            label = float(e[3])  # int(e[3])
+
+            feat = np.array([float(x) for x in e[4:]])
+
+            u_list.append(u)
+            i_list.append(i)
+            ts_list.append(ts)
+            label_list.append(label)
+            idx_list.append(idx)
+
+            feat_l.append(feat)
+    return pd.DataFrame({'u': u_list,
+                         'i': i_list,
+                         'ts': ts_list,
+                         'label': label_list,
+                         'idx': idx_list}), np.array(feat_l)
+
+# Re index nodes for DGL convience
+def reindex(df, bipartite=True):
+    new_df = df.copy()
+    if bipartite:
+        assert (df.u.max() - df.u.min() + 1 == len(df.u.unique()))
+        assert (df.i.max() - df.i.min() + 1 == len(df.i.unique()))
+
+        upper_u = df.u.max() + 1
+        new_i = df.i + upper_u
+
+        new_df.i = new_i
+        new_df.u += 1
+        new_df.i += 1
+        new_df.idx += 1
+    else:
+        new_df.u += 1
+        new_df.i += 1
+        new_df.idx += 1
+
+    return new_df
+
+# Save edge list, features in different file for data easy process data
+def run(data_name, bipartite=True):
+    PATH = './data/{}.csv'.format(data_name)
+    OUT_DF = './data/ml_{}.csv'.format(data_name)
+    OUT_FEAT = './data/ml_{}.npy'.format(data_name)
+    OUT_NODE_FEAT = './data/ml_{}_node.npy'.format(data_name)
+
+    df, feat = preprocess(PATH)
+    new_df = reindex(df, bipartite)
+
+    empty = np.zeros(feat.shape[1])[np.newaxis, :]
+    feat = np.vstack([empty, feat])
+
+    max_idx = max(new_df.u.max(), new_df.i.max())
+    rand_feat = np.zeros((max_idx + 1, 172))
+
+    new_df.to_csv(OUT_DF)
+    np.save(OUT_FEAT, feat)
+    np.save(OUT_NODE_FEAT, rand_feat)
+
+# === code from twitter-research-tgn end ===
+
+# If you have new dataset follow by same format in Jodie,
+# you can directly use name to retrieve dataset
+
+def TemporalDataset(dataset):
+    if not os.path.exists('./data/{}.bin'.format(dataset)):
+        if not os.path.exists('./data/{}.csv'.format(dataset)):
+            if not os.path.exists('./data'):
+                os.mkdir('./data')
+
+            url = 'https://snap.stanford.edu/jodie/{}.csv'.format(dataset)
+            print("Start Downloading File....")
+            context = ssl._create_unverified_context()
+            data = urllib.request.urlopen(url, context=context)
+            with open("./data/{}.csv".format(dataset), "wb") as handle:
+                handle.write(data.read())
+
+        print("Start Process Data ...")
+        run(dataset)
+        raw_connection = pd.read_csv('./data/ml_{}.csv'.format(dataset))
+        raw_feature = np.load('./data/ml_{}.npy'.format(dataset))
+        # -1 for re-index the node
+        src = raw_connection['u'].to_numpy()-1
+        dst = raw_connection['i'].to_numpy()-1
+        # Create directed graph
+        g = dgl.graph((src, dst))
+        g.edata['timestamp'] = torch.from_numpy(
+            raw_connection['ts'].to_numpy())
+        g.edata['label'] = torch.from_numpy(raw_connection['label'].to_numpy())
+        g.edata['feats'] = torch.from_numpy(raw_feature[1:, :]).float()
+        dgl.save_graphs('./data/{}.bin'.format(dataset), [g])
+    else:
+        print("Data is exist directly loaded.")
+        gs, _ = dgl.load_graphs('./data/{}.bin'.format(dataset))
+        g = gs[0]
+    return g
+
+def TemporalWikipediaDataset():
+    # Download the dataset
+    return TemporalDataset('wikipedia')
+
+def TemporalRedditDataset():
+    return TemporalDataset('reddit')

examples/pytorch/tgn/dataloading.py

+import torch
+import dgl
+
+from dgl.dataloading.dataloader import EdgeCollator, assign_block_eids
+from dgl.dataloading import BlockSampler
+from dgl.dataloading.pytorch import _pop_subgraph_storage, _pop_blocks_storage
+from dgl.base import DGLError
+
+from functools import partial
+import copy
+import dgl.function as fn
+
+
+def _prepare_tensor(g, data, name, is_distributed):
+    return torch.tensor(data) if is_distributed else dgl.utils.prepare_tensor(g, data, name)
+
+
+class TemporalSampler(BlockSampler):
+    """ Temporal Sampler builds computational and temporal dependency of node representations via
+    temporal neighbors selection and screening.
+
+    The sampler expects input node to have same time stamps, in the case of TGN, it should be 
+    either positive [src,dst] pair or negative samples. It will first take in-subgraph of seed
+    nodes and then screening out edges which happen after that timestamp. Finally it will sample
+    a fixed number of neighbor edges using random or topk sampling.
+
+    Parameters
+    ----------
+    sampler_type : str
+        sampler indication string of the final sampler.
+
+        If 'topk' then sample topk most recent nodes
+
+        If 'uniform' then uniform randomly sample k nodes
+
+    k : int
+        maximum number of neighors to sampler
+
+        default 10 neighbors as paper stated
+
+    Examples
+    ----------
+    Please refers to examples/pytorch/tgn/train.py
+
+    """
+
+    def __init__(self, sampler_type='topk', k=10):
+        super(TemporalSampler, self).__init__(1, False)
+        if sampler_type == 'topk':
+            self.sampler = partial(
+                dgl.sampling.select_topk, k=k, weight='timestamp')
+        elif sampler_type == 'uniform':
+            self.sampler = partial(dgl.sampling.sample_neighbors, fanout=k)
+        else:
+            raise DGLError(
+                "Sampler string invalid please use \'topk\' or \'uniform\'")
+
+    def sampler_frontier(self,
+                         block_id,
+                         g,
+                         seed_nodes,
+                         timestamp):
+        full_neighbor_subgraph = dgl.in_subgraph(g, seed_nodes)
+        full_neighbor_subgraph = dgl.add_edges(full_neighbor_subgraph,
+                                               seed_nodes, seed_nodes)
+
+        temporal_edge_mask = (full_neighbor_subgraph.edata['timestamp'] < timestamp) + (
+            full_neighbor_subgraph.edata['timestamp'] <= 0)
+        temporal_subgraph = dgl.edge_subgraph(
+            full_neighbor_subgraph, temporal_edge_mask)
+
+        # Map preserve ID
+        temp2origin = temporal_subgraph.ndata[dgl.NID]
+
+        # The added new edgge will be preserved hence
+        root2sub_dict = dict(
+            zip(temp2origin.tolist(), temporal_subgraph.nodes().tolist()))
+        temporal_subgraph.ndata[dgl.NID] = g.ndata[dgl.NID][temp2origin]
+        seed_nodes = [root2sub_dict[int(n)] for n in seed_nodes]
+        final_subgraph = self.sampler(g=temporal_subgraph, nodes=seed_nodes)
+        final_subgraph.remove_self_loop()
+        return final_subgraph
+
+        # Temporal Subgraph
+    def sample_blocks(self,
+                      g,
+                      seed_nodes,
+                      timestamp):
+        blocks = []
+        frontier = self.sampler_frontier(0, g, seed_nodes, timestamp)
+        #block = transform.to_block(frontier,seed_nodes)
+        block = frontier
+        if self.return_eids:
+            assign_block_eids(block, frontier)
+        blocks.append(block)
+        return blocks
+
+
+class TemporalEdgeCollator(EdgeCollator):
+    """ Temporal Edge collator merge the edges specified by eid: items
+
+    Since we cannot keep duplicated nodes on a graph we need to iterate though
+    the incoming edges and expand the duplicated node and form a batched block
+    graph capture the temporal and computational dependency.
+
+    Parameters
+    ----------
+
+    g : DGLGraph
+        The graph from which the edges are iterated in minibatches and the subgraphs
+        are generated.
+
+    eids : Tensor or dict[etype, Tensor]
+        The edge set in graph :attr:`g` to compute outputs.
+
+    block_sampler : dgl.dataloading.BlockSampler
+        The neighborhood sampler.
+
+    g_sampling : DGLGraph, optional
+        The graph where neighborhood sampling and message passing is performed.
+        Note that this is not necessarily the same as :attr:`g`.
+        If None, assume to be the same as :attr:`g`.
+
+    exclude : str, optional
+        Whether and how to exclude dependencies related to the sampled edges in the
+        minibatch.  Possible values are
+
+        * None, which excludes nothing.
+
+        * ``'reverse_id'``, which excludes the reverse edges of the sampled edges.  The said
+          reverse edges have the same edge type as the sampled edges.  Only works
+          on edge types whose source node type is the same as its destination node type.
+
+        * ``'reverse_types'``, which excludes the reverse edges of the sampled edges.  The
+          said reverse edges have different edge types from the sampled edges.
+
+        If ``g_sampling`` is given, ``exclude`` is ignored and will be always ``None``.
+
+    reverse_eids : Tensor or dict[etype, Tensor], optional
+        The mapping from original edge ID to its reverse edge ID.
+        Required and only used when ``exclude`` is set to ``reverse_id``.
+        For heterogeneous graph this will be a dict of edge type and edge IDs.  Note that
+        only the edge types whose source node type is the same as destination node type
+        are needed.
+
+    reverse_etypes : dict[etype, etype], optional
+        The mapping from the edge type to its reverse edge type.
+        Required and only used when ``exclude`` is set to ``reverse_types``.
+
+    negative_sampler : callable, optional
+        The negative sampler.  Can be omitted if no negative sampling is needed.
+        The negative sampler must be a callable that takes in the following arguments:
+
+        * The original (heterogeneous) graph.
+
+        * The ID array of sampled edges in the minibatch, or the dictionary of edge
+          types and ID array of sampled edges in the minibatch if the graph is
+          heterogeneous.
+
+        It should return
+
+        * A pair of source and destination node ID arrays as negative samples,
+          or a dictionary of edge types and such pairs if the graph is heterogenenous.
+
+        A set of builtin negative samplers are provided in
+        :ref:`the negative sampling module <api-dataloading-negative-sampling>`.
+
+    example
+    ----------
+    Please refers to examples/pytorch/tgn/train.py
+
+    """
+
+    def _collate_with_negative_sampling(self, items):
+        items = _prepare_tensor(self.g_sampling, items, 'items', False)
+        # Here node id will not change
+        pair_graph = self.g.edge_subgraph(items, preserve_nodes=True)
+        induced_edges = pair_graph.edata[dgl.EID]
+
+        neg_srcdst_raw = self.negative_sampler(self.g, items)
+        neg_srcdst = {self.g.canonical_etypes[0]: neg_srcdst_raw}
+        dtype = list(neg_srcdst.values())[0][0].dtype
+        neg_edges = {
+            etype: neg_srcdst.get(etype, (torch.tensor(
+                [], dtype=dtype), torch.tensor([], dtype=dtype)))
+            for etype in self.g.canonical_etypes}
+
+        neg_pair_graph = dgl.heterograph(
+            neg_edges, {ntype: self.g.number_of_nodes(ntype) for ntype in self.g.ntypes})
+        pair_graph, neg_pair_graph = dgl.transform.compact_graphs(
+            [pair_graph, neg_pair_graph])
+        # Need to remap id
+        pair_graph.ndata[dgl.NID] = self.g.nodes()[pair_graph.ndata[dgl.NID]]
+        neg_pair_graph.ndata[dgl.NID] = self.g.nodes()[
+            neg_pair_graph.ndata[dgl.NID]]
+
+        pair_graph.edata[dgl.EID] = induced_edges
+
+        batch_graphs = []
+        nodes_id = []
+        timestamps = []
+
+        for i, edge in enumerate(zip(self.g.edges()[0][items], self.g.edges()[1][items])):
+            ts = pair_graph.edata['timestamp'][i]
+            timestamps.append(ts)
+            subg = self.block_sampler.sample_blocks(self.g_sampling,
+                                                    list(edge),
+                                                    timestamp=ts)[0]
+            subg.ndata['timestamp'] = ts.repeat(subg.num_nodes())
+            nodes_id.append(subg.srcdata[dgl.NID])
+            batch_graphs.append(subg)
+        timestamps = torch.tensor(timestamps).repeat_interleave(
+            self.negative_sampler.k)
+        for i, neg_edge in enumerate(zip(neg_srcdst_raw[0].tolist(), neg_srcdst_raw[1].tolist())):
+            ts = timestamps[i]
+            subg = self.block_sampler.sample_blocks(self.g_sampling,
+                                                    [neg_edge[1]],
+                                                    timestamp=ts)[0]
+            subg.ndata['timestamp'] = ts.repeat(subg.num_nodes())
+            batch_graphs.append(subg)
+        blocks = [dgl.batch(batch_graphs)]
+        input_nodes = torch.cat(nodes_id)
+        return input_nodes, pair_graph, neg_pair_graph, blocks
+
+    def collator(self, items):
+        """
+        The interface of collator, input items is edge id of the attached graph
+        """
+        result = super().collate(items)
+        # Copy the feature from parent graph
+        _pop_subgraph_storage(result[1], self.g)
+        _pop_subgraph_storage(result[2], self.g)
+        _pop_blocks_storage(result[-1], self.g_sampling)
+        return result
+
+
+class TemporalEdgeDataLoader(dgl.dataloading.EdgeDataLoader):
+    """ TemporalEdgeDataLoader is an iteratable object to generate blocks for temporal embedding
+    as well as pos and neg pair graph for memory update.
+
+    The batch generated will follow temporal order
+
+    Parameters
+    ----------
+    g : dgl.Heterograph
+        graph for batching the temporal edge id as well as generate negative subgraph
+
+    eids : torch.tensor() or numpy array
+        eids range which to be batched, it is useful to split training validation test dataset
+
+    block_sampler : dgl.dataloading.BlockSampler
+        temporal neighbor sampler which sample temporal and computationally depend blocks for computation
+
+    device : str
+        'cpu' means load dataset on cpu
+        'cuda' means load dataset on gpu
+
+    collator : dgl.dataloading.EdgeCollator
+        Merge input eid from pytorch dataloader to graph
+
+    Example
+    ----------
+    Please refers to examples/pytorch/tgn/train.py
+
+    """
+
+    def __init__(self, g, eids, block_sampler, device='cpu', collator=TemporalEdgeCollator, **kwargs):
+        collator_kwargs = {}
+        dataloader_kwargs = {}
+        for k, v in kwargs.items():
+            if k in self.collator_arglist:
+                collator_kwargs[k] = v
+            else:
+                dataloader_kwargs[k] = v
+        self.collator = collator(g, eids, block_sampler, **collator_kwargs)
+
+        assert not isinstance(g, dgl.distributed.DistGraph), \
+            'EdgeDataLoader does not support DistGraph for now. ' \
+            + 'Please use DistDataLoader directly.'
+        self.dataloader = torch.utils.data.DataLoader(
+            self.collator.dataset, collate_fn=self.collator.collate, **dataloader_kwargs)
+        self.device = device
+
+        # Precompute the CSR and CSC representations so each subprocess does not
+        # duplicate.
+        if dataloader_kwargs.get('num_workers', 0) > 0:
+            g.create_formats_()
+
+# ====== Fast Mode ======
+
+# Part of code in reservoir sampling comes from PyG library
+# https://github.com/rusty1s/pytorch_geometric/nn/models/tgn.py
+
+
+class FastTemporalSampler(BlockSampler):
+    """Temporal Sampler which implemented with a fast query lookup table. Sample
+    temporal and computationally depending subgraph.
+
+    The sampler maintains a lookup table of most current k neighbors of each node
+    each time, the sampler need to be updated with new edges from incoming batch to
+    update the lookup table.
+
+    Parameters
+    ----------
+    g : dgl.Heterograph
+        graph to be sampled here it which only exist to provide feature and data reference
+
+    k : int
+        number of neighbors the lookup table is maintaining
+
+    device : str
+        indication str which represent where the data will be stored
+        'cpu' store the intermediate data on cpu memory
+        'cuda' store the intermediate data on gpu memory 
+
+    Example
+    ----------
+    Please refers to examples/pytorch/tgn/train.py
+    """
+
+    def __init__(self, g, k, device='cpu'):
+        self.k = k
+        self.g = g
+        num_nodes = g.num_nodes()
+        self.neighbors = torch.empty(
+            (num_nodes, k), dtype=torch.long, device=device)
+        self.e_id = torch.empty(
+            (num_nodes, k), dtype=torch.long, device=device)
+        self.__assoc__ = torch.empty(
+            num_nodes, dtype=torch.long, device=device)
+        self.last_update = torch.zeros(num_nodes, dtype=torch.double)
+        self.reset()
+
+    def sample_frontier(self,
+                        block_id,
+                        g,
+                        seed_nodes):
+        n_id = seed_nodes
+        # Here Assume n_id is the bg nid
+        neighbors = self.neighbors[n_id]
+        nodes = n_id.view(-1, 1).repeat(1, self.k)
+        e_id = self.e_id[n_id]
+        mask = e_id >= 0
+
+        neighbors[~mask] = nodes[~mask]
+        # Screen out orphan node
+
+        orphans = nodes[~mask].unique()
+        nodes = nodes[mask]
+        neighbors = neighbors[mask]
+
+        e_id = e_id[mask]
+        neighbors = neighbors.flatten()
+        nodes = nodes.flatten()
+        n_id = torch.cat([nodes, neighbors]).unique()
+        self.__assoc__[n_id] = torch.arange(n_id.size(0), device=n_id.device)
+
+        neighbors, nodes = self.__assoc__[neighbors], self.__assoc__[nodes]
+        subg = dgl.graph((nodes, neighbors))
+
+        # New node to complement orphans which haven't created
+        subg.add_nodes(len(orphans))
+
+        # Copy the seed node feature to subgraph
+        subg.edata['timestamp'] = torch.zeros(subg.num_edges()).double()
+        subg.edata['timestamp'] = self.g.edata['timestamp'][e_id]
+
+        n_id = torch.cat([n_id, orphans])
+        subg.ndata['timestamp'] = self.last_update[n_id]
+        subg.edata['feats'] = torch.zeros(
+            (subg.num_edges(), self.g.edata['feats'].shape[1])).float()
+        subg.edata['feats'] = self.g.edata['feats'][e_id]
+        subg = dgl.add_self_loop(subg)
+        subg.ndata[dgl.NID] = n_id
+        return subg
+
+    def sample_blocks(self,
+                      g,
+                      seed_nodes):
+        blocks = []
+        frontier = self.sample_frontier(0, g, seed_nodes)
+        block = frontier
+        blocks.append(block)
+        return blocks
+
+    def add_edges(self, src, dst):
+        """
+        Add incoming batch edge info to the lookup table
+
+        Parameters
+        ----------
+        src : torch.Tensor
+            src node of incoming batch of it should be consistent with self.g
+
+        dst : torch.Tensor
+            src node of incoming batch of it should be consistent with self.g
+        """
+        neighbors = torch.cat([src, dst], dim=0)
+        nodes = torch.cat([dst, src], dim=0)
+        e_id = torch.arange(self.cur_e_id, self.cur_e_id + src.size(0),
+                            device=src.device).repeat(2)
+        self.cur_e_id += src.numel()
+
+        # Convert newly encountered interaction ids so that they point to
+        # locations of a "dense" format of shape [num_nodes, size].
+        nodes, perm = nodes.sort()
+        neighbors, e_id = neighbors[perm], e_id[perm]
+
+        n_id = nodes.unique()
+        self.__assoc__[n_id] = torch.arange(n_id.numel(), device=n_id.device)
+
+        dense_id = torch.arange(nodes.size(0), device=nodes.device) % self.k
+        dense_id += self.__assoc__[nodes].mul_(self.k)
+
+        dense_e_id = e_id.new_full((n_id.numel() * self.k, ), -1)
+        dense_e_id[dense_id] = e_id
+        dense_e_id = dense_e_id.view(-1, self.k)
+
+        dense_neighbors = e_id.new_empty(n_id.numel() * self.k)
+        dense_neighbors[dense_id] = neighbors
+        dense_neighbors = dense_neighbors.view(-1, self.k)
+
+        # Collect new and old interactions...
+        e_id = torch.cat([self.e_id[n_id, :self.k], dense_e_id], dim=-1)
+        neighbors = torch.cat(
+            [self.neighbors[n_id, :self.k], dense_neighbors], dim=-1)
+
+        # And sort them based on `e_id`.
+        e_id, perm = e_id.topk(self.k, dim=-1)
+        self.e_id[n_id] = e_id
+        self.neighbors[n_id] = torch.gather(neighbors, 1, perm)
+
+    def reset(self):
+        """
+        Clean up the lookup table
+        """
+        self.cur_e_id = 0
+        self.e_id.fill_(-1)
+
+    def attach_last_update(self, last_t):
+        """
+        Attach current last timestamp a node has been updated
+
+        Parameters:
+        ----------
+        last_t : torch.Tensor
+            last timestamp a node has been updated its size need to be consistent with self.g
+
+        """
+        self.last_update = last_t
+
+    def sync(self, sampler):
+        """
+        Copy the lookup table information from another sampler
+
+        This method is useful run the test dataset with new node,
+        when test new node dataset the lookup table's state should
+        be restored from the sampler just after validation 
+
+        Parameters
+        ----------
+        sampler : FastTemporalSampler
+            The sampler from which current sampler get the lookup table info
+        """
+        self.cur_e_id = sampler.cur_e_id
+        self.neighbors = copy.deepcopy(sampler.neighbors)
+        self.e_id = copy.deepcopy(sampler.e_id)
+        self.__assoc__ = copy.deepcopy(sampler.__assoc__)
+
+
+class FastTemporalEdgeCollator(EdgeCollator):
+    """ Temporal Edge collator merge the edges specified by eid: items
+
+    Since we cannot keep duplicated nodes on a graph we need to iterate though
+    the incoming edges and expand the duplicated node and form a batched block
+    graph capture the temporal and computational dependency.
+
+    Parameters
+    ----------
+
+    g : DGLGraph
+        The graph from which the edges are iterated in minibatches and the subgraphs
+        are generated.
+
+    eids : Tensor or dict[etype, Tensor]
+        The edge set in graph :attr:`g` to compute outputs.
+
+    block_sampler : dgl.dataloading.BlockSampler
+        The neighborhood sampler.
+
+    g_sampling : DGLGraph, optional
+        The graph where neighborhood sampling and message passing is performed.
+        Note that this is not necessarily the same as :attr:`g`.
+        If None, assume to be the same as :attr:`g`.
+
+    exclude : str, optional
+        Whether and how to exclude dependencies related to the sampled edges in the
+        minibatch.  Possible values are
+
+        * None, which excludes nothing.
+
+        * ``'reverse_id'``, which excludes the reverse edges of the sampled edges.  The said
+          reverse edges have the same edge type as the sampled edges.  Only works
+          on edge types whose source node type is the same as its destination node type.
+
+        * ``'reverse_types'``, which excludes the reverse edges of the sampled edges.  The
+          said reverse edges have different edge types from the sampled edges.
+
+        If ``g_sampling`` is given, ``exclude`` is ignored and will be always ``None``.
+
+    reverse_eids : Tensor or dict[etype, Tensor], optional
+        The mapping from original edge ID to its reverse edge ID.
+        Required and only used when ``exclude`` is set to ``reverse_id``.
+        For heterogeneous graph this will be a dict of edge type and edge IDs.  Note that
+        only the edge types whose source node type is the same as destination node type
+        are needed.
+
+    reverse_etypes : dict[etype, etype], optional
+        The mapping from the edge type to its reverse edge type.
+        Required and only used when ``exclude`` is set to ``reverse_types``.
+
+    negative_sampler : callable, optional
+        The negative sampler.  Can be omitted if no negative sampling is needed.
+        The negative sampler must be a callable that takes in the following arguments:
+
+        * The original (heterogeneous) graph.
+
+        * The ID array of sampled edges in the minibatch, or the dictionary of edge
+          types and ID array of sampled edges in the minibatch if the graph is
+          heterogeneous.
+
+        It should return
+
+        * A pair of source and destination node ID arrays as negative samples,
+          or a dictionary of edge types and such pairs if the graph is heterogenenous.
+
+        A set of builtin negative samplers are provided in
+        :ref:`the negative sampling module <api-dataloading-negative-sampling>`.
+
+    example
+    ----------
+    Please refers to examples/pytorch/tgn/train.py
+
+    """
+
+    def _collate_with_negative_sampling(self, items):
+        items = _prepare_tensor(self.g_sampling, items, 'items', False)
+        # Here node id will not change
+        pair_graph = self.g.edge_subgraph(items, preserve_nodes=True)
+        induced_edges = pair_graph.edata[dgl.EID]
+
+        neg_srcdst_raw = self.negative_sampler(self.g, items)
+        neg_srcdst = {self.g.canonical_etypes[0]: neg_srcdst_raw}
+        dtype = list(neg_srcdst.values())[0][0].dtype
+        neg_edges = {
+            etype: neg_srcdst.get(etype, (torch.tensor(
+                [], dtype=dtype), torch.tensor([], dtype=dtype)))
+            for etype in self.g.canonical_etypes}
+
+        neg_pair_graph = dgl.heterograph(
+            neg_edges, {ntype: self.g.number_of_nodes(ntype) for ntype in self.g.ntypes})
+        pair_graph, neg_pair_graph = dgl.transform.compact_graphs(
+            [pair_graph, neg_pair_graph])
+        # Need to remap id
+
+        pair_graph.ndata[dgl.NID] = self.g.nodes()[pair_graph.ndata[dgl.NID]]
+        neg_pair_graph.ndata[dgl.NID] = self.g.nodes()[
+            neg_pair_graph.ndata[dgl.NID]]
+
+        pair_graph.edata[dgl.EID] = induced_edges
+
+        seed_nodes = pair_graph.ndata[dgl.NID]
+        blocks = self.block_sampler.sample_blocks(self.g_sampling, seed_nodes)
+        blocks[0].ndata['timestamp'] = torch.zeros(
+            blocks[0].num_nodes()).double()
+        input_nodes = blocks[0].edges()[1]
+
+        # update sampler
+        _src = self.g.nodes()[self.g.edges()[0][items]]
+        _dst = self.g.nodes()[self.g.edges()[1][items]]
+        self.block_sampler.add_edges(_src, _dst)
+        return input_nodes, pair_graph, neg_pair_graph, blocks
+
+    def collator(self, items):
+        result = super().collate(items)
+        # Copy the feature from parent graph
+        _pop_subgraph_storage(result[1], self.g)
+        _pop_subgraph_storage(result[2], self.g)
+        _pop_blocks_storage(result[-1], self.g_sampling)
+        return result
+
+
+# ====== Simple Mode ======
+
+# Part of code comes from paper
+# "APAN: Asynchronous Propagation Attention Network for Real-time Temporal Graph Embedding"
+# that will be appeared in SIGMOD 21, code repo https://github.com/WangXuhongCN/APAN
+
+class SimpleTemporalSampler(dgl.dataloading.BlockSampler):
+    '''
+    Simple Temporal Sampler just choose the edges that happen before the current timestamp, to build the subgraph of the corresponding nodes. 
+    And then the sampler uses the simplest static graph neighborhood sampling methods.
+
+    Parameters
+    ----------
+
+    fanouts : [int, ..., int] int list
+        The neighbors sampling strategy 
+
+    '''
+
+    def __init__(self, g, fanouts, return_eids=False):
+        super().__init__(len(fanouts), return_eids)
+
+        self.fanouts = fanouts  
+        self.ts = 0
+        self.frontiers = [None for _ in range(len(fanouts))]
+
+    def sample_frontier(self, block_id, g, seed_nodes):
+        '''
+        Deleting the the edges that happen after the current timestamp, then use a simple topk edge sampling by timestamp.
+        '''
+        fanout = self.fanouts[block_id]
+        # List of neighbors to sample per edge type for each GNN layer, starting from the first layer.
+        g = dgl.in_subgraph(g, seed_nodes)  
+        g.remove_edges(torch.where(g.edata['timestamp'] > self.ts)[0])  # Deleting the the edges that happen after the current timestamp
+
+        if fanout is None:  # full neighborhood sampling
+            frontier = g
+        else:
+            frontier = dgl.sampling.select_topk(g, fanout, 'timestamp', seed_nodes)  # most recent timestamp edge sampling
+        self.frontiers[block_id] = frontier  # save frontier
+        return frontier
+
+
+class SimpleTemporalEdgeCollator(dgl.dataloading.EdgeCollator):
+    '''
+    Temporal Edge collator merge the edges specified by eid: items
+
+    
+
+    Parameters
+    ----------
+
+    g : DGLGraph
+        The graph from which the edges are iterated in minibatches and the subgraphs
+        are generated.
+
+    eids : Tensor or dict[etype, Tensor]
+        The edge set in graph :attr:`g` to compute outputs.
+
+    block_sampler : dgl.dataloading.BlockSampler
+        The neighborhood sampler.
+
+    g_sampling : DGLGraph, optional
+        The graph where neighborhood sampling and message passing is performed.
+        Note that this is not necessarily the same as :attr:`g`.
+        If None, assume to be the same as :attr:`g`.
+
+    exclude : str, optional
+        Whether and how to exclude dependencies related to the sampled edges in the
+        minibatch.  Possible values are
+
+        * None, which excludes nothing.
+
+        * ``'reverse_id'``, which excludes the reverse edges of the sampled edges.  The said
+          reverse edges have the same edge type as the sampled edges.  Only works
+          on edge types whose source node type is the same as its destination node type.
+
+        * ``'reverse_types'``, which excludes the reverse edges of the sampled edges.  The
+          said reverse edges have different edge types from the sampled edges.
+
+        If ``g_sampling`` is given, ``exclude`` is ignored and will be always ``None``.
+
+    reverse_eids : Tensor or dict[etype, Tensor], optional
+        The mapping from original edge ID to its reverse edge ID.
+        Required and only used when ``exclude`` is set to ``reverse_id``.
+        For heterogeneous graph this will be a dict of edge type and edge IDs.  Note that
+        only the edge types whose source node type is the same as destination node type
+        are needed.
+
+    reverse_etypes : dict[etype, etype], optional
+        The mapping from the edge type to its reverse edge type.
+        Required and only used when ``exclude`` is set to ``reverse_types``.
+
+    negative_sampler : callable, optional
+        The negative sampler.  Can be omitted if no negative sampling is needed.
+        The negative sampler must be a callable that takes in the following arguments:
+
+        * The original (heterogeneous) graph.
+
+        * The ID array of sampled edges in the minibatch, or the dictionary of edge
+          types and ID array of sampled edges in the minibatch if the graph is
+          heterogeneous.
+
+        It should return
+
+        * A pair of source and destination node ID arrays as negative samples,
+          or a dictionary of edge types and such pairs if the graph is heterogenenous.
+
+        A set of builtin negative samplers are provided in
+        :ref:`the negative sampling module <api-dataloading-negative-sampling>`.
+    '''
+    def __init__(self, g, eids, block_sampler, g_sampling=None, exclude=None,
+                reverse_eids=None, reverse_etypes=None, negative_sampler=None):
+        super(SimpleTemporalEdgeCollator,self).__init__(g,eids,block_sampler,
+                                                 g_sampling,exclude,reverse_eids,reverse_etypes,negative_sampler)
+        self.n_layer = len(self.block_sampler.fanouts)
+
+    def collate(self,items): 
+        '''
+        items: edge id in graph g.
+        We sample iteratively k-times and batch them into one single subgraph.
+        '''
+        current_ts = self.g.edata['timestamp'][items[0]]     #only sample edges before current timestamp
+        self.block_sampler.ts = current_ts    # restore the current timestamp to the graph sampler.
+
+        # if link prefiction, we use a negative_sampler to generate neg-graph for loss computing.
+        if self.negative_sampler is None:
+            neg_pair_graph = None
+            input_nodes, pair_graph, blocks = self._collate(items)
+        else:
+            input_nodes, pair_graph, neg_pair_graph, blocks = self._collate_with_negative_sampling(items)
+
+        # we sampling k-hop subgraph and batch them into one graph
+        for i in range(self.n_layer-1):
+            self.block_sampler.frontiers[0].add_edges(*self.block_sampler.frontiers[i+1].edges())
+        frontier = self.block_sampler.frontiers[0]
+        # computing node last-update timestamp
+        frontier.update_all(fn.copy_e('timestamp','ts'), fn.max('ts','timestamp'))
+    
+        return input_nodes, pair_graph, neg_pair_graph, [frontier]
+
+
+
+

examples/pytorch/tgn/modules.py

+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import dgl
+from dgl.base import DGLError
+from dgl.ops import edge_softmax
+import dgl.function as fn
+
+
+class MergeLayer(nn.Module):
+    """Merge two tensor into one
+    Which is useful as skip connection in Merge GAT's input with output
+
+    Parameter
+    ----------
+    dim1 : int
+        dimension of first input tensor
+
+    dim2 : int
+        dimension of second input tensor
+
+    dim3 : int
+        hidden dimension after first merging
+
+    dim4 : int
+        output dimension
+
+    Example
+    ----------
+    >>> merger = MergeLayer(10,10,10,5)
+    >>> input1 = torch.ones(4,10)
+    >>> input2 = torch.ones(4,10)
+    >>> merger(input1,input2)
+    tensor([[-0.1578,  0.1842,  0.2373,  1.2656,  1.0362],
+        [-0.1578,  0.1842,  0.2373,  1.2656,  1.0362],
+        [-0.1578,  0.1842,  0.2373,  1.2656,  1.0362],
+        [-0.1578,  0.1842,  0.2373,  1.2656,  1.0362]],
+       grad_fn=<AddmmBackward>)
+    """
+
+    def __init__(self, dim1, dim2, dim3, dim4):
+        super(MergeLayer, self).__init__()
+        self.fc1 = torch.nn.Linear(dim1 + dim2, dim3)
+        self.fc2 = torch.nn.Linear(dim3, dim4)
+        self.act = torch.nn.ReLU()
+
+        torch.nn.init.xavier_normal_(self.fc1.weight)
+        torch.nn.init.xavier_normal_(self.fc2.weight)
+
+    def forward(self, x1, x2):
+        x = torch.cat([x1, x2], dim=1)
+        h = self.act(self.fc1(x))
+        return self.fc2(h)
+
+
+class MsgLinkPredictor(nn.Module):
+    """Predict Pair wise link from pos subg and neg subg
+    use message passing.
+
+    Use Two layer MLP on edge to predict the link probability
+
+    Parameters
+    ----------
+    embed_dim : int
+        dimension of each each feature's embedding
+
+    Example
+    ----------
+    >>> linkpred = MsgLinkPredictor(10)
+    >>> pos_g = dgl.graph(([0,1,2,3,4],[1,2,3,4,0]))
+    >>> neg_g = dgl.graph(([0,1,2,3,4],[2,1,4,3,0]))
+    >>> x = torch.ones(5,10)
+    >>> linkpred(x,pos_g,neg_g)
+    (tensor([[0.0902],
+         [0.0902],
+         [0.0902],
+         [0.0902],
+         [0.0902]], grad_fn=<AddmmBackward>),
+    tensor([[0.0902],
+         [0.0902],
+         [0.0902],
+         [0.0902],
+         [0.0902]], grad_fn=<AddmmBackward>))
+    """
+
+    def __init__(self, emb_dim):
+        super(MsgLinkPredictor, self).__init__()
+        self.src_fc = nn.Linear(emb_dim, emb_dim)
+        self.dst_fc = nn.Linear(emb_dim, emb_dim)
+        self.out_fc = nn.Linear(emb_dim, 1)
+
+    def link_pred(self, edges):
+        src_hid = self.src_fc(edges.src['embedding'])
+        dst_hid = self.dst_fc(edges.dst['embedding'])
+        score = F.relu(src_hid+dst_hid)
+        score = self.out_fc(score)
+        return {'score': score}
+
+    def forward(self, x, pos_g, neg_g):
+        # Local Scope?
+        pos_g.ndata['embedding'] = x
+        neg_g.ndata['embedding'] = x
+
+        pos_g.apply_edges(self.link_pred)
+        neg_g.apply_edges(self.link_pred)
+
+        pos_escore = pos_g.edata['score']
+        neg_escore = neg_g.edata['score']
+        return pos_escore, neg_escore
+
+
+class TimeEncode(nn.Module):
+    """Use finite fourier series with different phase and frequency to encode
+    time different between two event
+
+    ..math::
+        \Phi(t) = [\cos(\omega_0t+\psi_0),\cos(\omega_1t+\psi_1),...,\cos(\omega_nt+\psi_n)] 
+
+    Parameter
+    ----------
+    dimension : int
+        Length of the fourier series. The longer it is , 
+        the more timescale information it can capture
+
+    Example
+    ----------
+    >>> tecd = TimeEncode(10)
+    >>> t = torch.tensor([[1]])
+    >>> tecd(t)
+    tensor([[[0.5403, 0.9950, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
+          1.0000, 1.0000]]], dtype=torch.float64, grad_fn=<CosBackward>)
+    """
+
+    def __init__(self, dimension):
+        super(TimeEncode, self).__init__()
+
+        self.dimension = dimension
+        self.w = torch.nn.Linear(1, dimension)
+
+        self.w.weight = torch.nn.Parameter((torch.from_numpy(1 / 10 ** np.linspace(0, 9, dimension)))
+                                           .double().reshape(dimension, -1))
+        self.w.bias = torch.nn.Parameter(torch.zeros(dimension).double())
+
+    def forward(self, t):
+        t = t.unsqueeze(dim=2)
+        output = torch.cos(self.w(t))
+        return output
+
+
+class MemoryModule(nn.Module):
+    """Memory module as well as update interface
+
+    The memory module stores both historical representation in last_update_t
+
+    Parameters
+    ----------
+    n_node : int
+        number of node of the entire graph
+
+    hidden_dim : int
+        dimension of memory of each node
+
+    Example
+    ----------
+    Please refers to examples/pytorch/tgn/tgn.py;
+                     examples/pytorch/tgn/train.py 
+
+    """
+
+    def __init__(self, n_node, hidden_dim):
+        super(MemoryModule, self).__init__()
+        self.n_node = n_node
+        self.hidden_dim = hidden_dim
+        self.reset_memory()
+
+    def reset_memory(self):
+        self.last_update_t = nn.Parameter(torch.zeros(
+            self.n_node).float(), requires_grad=False)
+        self.memory = nn.Parameter(torch.zeros(
+            (self.n_node, self.hidden_dim)).float(), requires_grad=False)
+
+    def backup_memory(self):
+        """
+        Return a deep copy of memory state and last_update_t
+        For test new node, since new node need to use memory upto validation set
+        After validation, memory need to be backed up before run test set without new node
+        so finally, we can use backup memory to update the new node test set
+        """
+        return self.memory.clone(), self.last_update_t.clone()
+
+    def restore_memory(self, memory_backup):
+        """Restore the memory from validation set
+
+        Parameters
+        ----------
+        memory_backup : (memory,last_update_t)
+            restore memory based on input tuple
+        """
+        self.memory = memory_backup[0].clone()
+        self.last_update_t = memory_backup[1].clone()
+
+    # Which is used for attach to subgraph
+    def get_memory(self, node_idxs):
+        return self.memory[node_idxs, :]
+
+    # When the memory need to be updated
+    def set_memory(self, node_idxs, values):
+        self.memory[node_idxs, :] = values
+
+    def set_last_update_t(self, node_idxs, values):
+        self.last_update_t[node_idxs] = values
+
+    # For safety check
+    def get_last_update(self, node_idxs):
+        return self.last_update_t[node_idxs]
+
+    def detach_memory(self):
+        """
+        Disconnect the memory from computation graph to prevent gradient be propagated multiple
+        times
+        """
+        self.memory.detach_()
+
+
+class MemoryOperation(nn.Module):
+    """ Memory update using message passing manner, update memory based on positive
+    pair graph of each batch with recurrent module GRU or RNN
+
+    Message function
+    ..math::
+        m_i(t) = concat(memory_i(t^-),TimeEncode(t),v_i(t))
+
+    v_i is node feature at current time stamp
+
+    Aggregation function
+    ..math::
+        \bar{m}_i(t) = last(m_i(t_1),...,m_i(t_b))
+
+    Update function
+    ..math::
+        memory_i(t) = GRU(\bar{m}_i(t),memory_i(t-1))
+
+    Parameters
+    ----------
+
+    updater_type : str
+        indicator string to specify updater
+
+        'rnn' : use Vanilla RNN as updater
+
+        'gru' : use GRU as updater
+
+    memory : MemoryModule
+        memory content for update
+
+    e_feat_dim : int
+        dimension of edge feature
+
+    temporal_dim : int
+        length of fourier series for time encoding
+
+    Example
+    ----------
+    Please refers to examples/pytorch/tgn/tgn.py
+    """
+
+    def __init__(self, updater_type, memory, e_feat_dim, temporal_dim):
+        super(MemoryOperation, self).__init__()
+        updater_dict = {'gru': nn.GRUCell, 'rnn': nn.RNNCell}
+        self.memory = memory
+        memory_dim = self.memory.hidden_dim
+        self.message_dim = memory_dim+memory_dim+e_feat_dim+temporal_dim
+        self.updater = updater_dict[updater_type](input_size=self.message_dim,
+                                                  hidden_size=memory_dim)
+        self.memory = memory
+        self.temporal_encoder = TimeEncode(temporal_dim)
+
+    # Here assume g is a subgraph from each iteration
+    def stick_feat_to_graph(self, g):
+        # How can I ensure order of the node ID
+        g.ndata['timestamp'] = self.memory.last_update_t[g.ndata[dgl.NID]]
+        g.ndata['memory'] = self.memory.memory[g.ndata[dgl.NID]]
+
+    def msg_fn_cat(self, edges):
+        src_delta_time = edges.data['timestamp'] - edges.src['timestamp']
+        time_encode = self.temporal_encoder(src_delta_time.unsqueeze(
+            dim=1)).view(len(edges.data['timestamp']), -1)
+        ret = torch.cat([edges.src['memory'], edges.dst['memory'],
+                         edges.data['feats'], time_encode], dim=1)
+        return {'message': ret, 'timestamp': edges.data['timestamp']}
+
+    def agg_last(self, nodes):
+        timestamp, latest_idx = torch.max(nodes.mailbox['timestamp'], dim=1)
+        ret = nodes.mailbox['message'].gather(1, latest_idx.repeat(
+            self.message_dim).view(-1, 1, self.message_dim)).view(-1, self.message_dim)
+        return {'message_bar': ret.reshape(-1, self.message_dim), 'timestamp': timestamp}
+
+    def update_memory(self, nodes):
+        # It should pass the feature through RNN
+        ret = self.updater(
+            nodes.data['message_bar'].float(), nodes.data['memory'].float())
+        return {'memory': ret}
+
+    def forward(self, g):
+        self.stick_feat_to_graph(g)
+        g.update_all(self.msg_fn_cat, self.agg_last, self.update_memory)
+        return g
+
+
+class TemporalGATConv(nn.Module):
+    """Dot Product based embedding with temporal encoding
+    it will each node will compute the attention weight with other
+    nodes using other node's memory as well as edge features
+
+    Aggregation:
+    ..math::
+        h_i^{(l+1)} = ReLu(\sum_{j\in \mathcal{N}(i)} \alpha_{i, j} (h_j^{(l)}(t)||e_{jj}||TimeEncode(t-t_j)))
+
+        \alpha_{i, j} = \mathrm{softmax_i}(\frac{QK^T}{\sqrt{d_k}})V
+
+    K,Q,V computation:
+    ..math::
+        K = W_k[memory_{src}(t),memory_{dst}(t),TimeEncode(t_{dst}-t_{src})]
+        Q = W_q[memory_{src}(t),memory_{dst}(t),TimeEncode(t_{dst}-t_{src})]
+        V = W_v[memory_{src}(t),memory_{dst}(t),TimeEncode(t_{dst}-t_{src})]
+
+    Parameters
+    ----------
+
+    edge_feats : int
+        dimension of edge feats
+
+    memory_feats : int
+        dimension of memory feats
+
+    temporal_feats : int
+        length of fourier series of time encoding
+
+    num_heads : int
+        number of head in multihead attention
+
+    allow_zero_in_degree : bool
+        Whether allow some node have indegree == 0 to prevent silence evaluation
+
+    Example
+    ----------
+    >>> attn = TemporalGATConv(2,2,5,2,2,False)
+    >>> star_graph = dgl.graph(([1,2,3,4,5],[0,0,0,0,0]))
+    >>> star_graph.edata['feats'] = torch.ones(5,2).double()
+    >>> star_graph.edata['timestamp'] = torch.zeros(5).double()
+    >>> memory = torch.ones(6,2)
+    >>> ts = torch.random.rand(6).double()
+    >>> star_graph = dgl.add_self_loop(star_graph)
+    >>> attn(graph,memory,ts)
+    tensor([[-0.0924, -0.3842],
+        [-0.0840, -0.3539],
+        [-0.0842, -0.3543],
+        [-0.0838, -0.3536],
+        [-0.0856, -0.3568],
+        [-0.0858, -0.3572]], grad_fn=<AddmmBackward>)
+    """
+
+    def __init__(self,
+                 edge_feats,
+                 memory_feats,
+                 temporal_feats,
+                 out_feats,
+                 num_heads,
+                 allow_zero_in_degree=False):
+        super(TemporalGATConv, self).__init__()
+        self._edge_feats = edge_feats
+        self._memory_feats = memory_feats
+        self._temporal_feats = temporal_feats
+        self._out_feats = out_feats
+        self._allow_zero_in_degree = allow_zero_in_degree
+        self._num_heads = num_heads
+
+        self.fc_Q = nn.Linear(self._memory_feats+self._temporal_feats,
+                              self._out_feats*self._num_heads, bias=False)
+        self.fc_K = nn.Linear(self._memory_feats+self._edge_feats +
+                              self._temporal_feats, self._out_feats*self._num_heads, bias=False)
+        self.fc_V = nn.Linear(self._memory_feats+self._edge_feats +
+                              self._temporal_feats, self._out_feats*self._num_heads, bias=False)
+        self.merge = MergeLayer(
+            self._memory_feats, self._out_feats*self._num_heads, 512, self._out_feats)
+        self.temporal_encoder = TimeEncode(self._temporal_feats)
+
+    def weight_fn(self, edges):
+        t0 = torch.zeros_like(edges.dst['timestamp'])
+        q = torch.cat([edges.dst['s'],
+                       self.temporal_encoder(t0.unsqueeze(dim=1)).view(len(t0), -1)], dim=1)
+        time_diff = edges.data['timestamp']-edges.src['timestamp']
+        time_encode = self.temporal_encoder(
+            time_diff.unsqueeze(dim=1)).view(len(t0), -1)
+        k = torch.cat(
+            [edges.src['s'], edges.data['feats'], time_encode], dim=1)
+        squeezed_k = self.fc_K(
+            k.float()).view(-1, self._num_heads, self._out_feats)
+        squeezed_q = self.fc_Q(
+            q.float()).view(-1, self._num_heads, self._out_feats)
+        ret = torch.sum(squeezed_q*squeezed_k, dim=2)
+        return {'a': ret, 'efeat': squeezed_k}
+
+    def msg_fn(self, edges):
+        ret = edges.data['sa'].view(-1, self._num_heads, 1)*edges.data['efeat']
+        return {'attn': ret}
+
+    def forward(self, graph, memory, ts):
+        graph = graph.local_var()  # Using local scope for graph
+        if not self._allow_zero_in_degree:
+            if(graph.in_degrees() == 0).any():
+                raise DGLError('There are 0-in-degree nodes in the graph, '
+                               'output for those nodes will be invalid. '
+                               'This is harmful for some applications, '
+                               'causing silent performance regression. '
+                               'Adding self-loop on the input graph by '
+                               'calling `g = dgl.add_self_loop(g)` will resolve '
+                               'the issue. Setting ``allow_zero_in_degree`` '
+                               'to be `True` when constructing this module will '
+                               'suppress the check and let the code run.')
+
+        #print("Shape: ",memory.shape,ts.shape)
+        graph.srcdata.update({'s': memory, 'timestamp': ts})
+        graph.dstdata.update({'s': memory, 'timestamp': ts})
+
+        # Dot product Calculate the attentio weight
+        graph.apply_edges(self.weight_fn)
+
+        # Edge softmax
+        graph.edata['sa'] = edge_softmax(
+            graph, graph.edata['a'])/(self._out_feats**0.5)
+
+        # Update dst node Here msg_fn include edge feature
+        graph.update_all(self.msg_fn, fn.sum('attn', 'agg_u'))
+
+        rst = graph.dstdata['agg_u']
+        # Implement skip connection
+        rst = self.merge(rst.view(-1, self._num_heads *
+                                  self._out_feats), graph.dstdata['s'])
+        return rst

examples/pytorch/tgn/tgn.py

+import copy
+import torch.nn as nn
+import dgl
+from modules import MemoryModule, MemoryOperation, TemporalGATConv, MsgLinkPredictor
+
+class TGN(nn.Module):
+    def __init__(self,
+                 edge_feat_dim,
+                 memory_dim,
+                 temporal_dim,
+                 embedding_dim,
+                 num_heads,
+                 num_nodes,  # entire graph
+                 n_neighbors=10,
+                 memory_updater_type='gru'):
+        super(TGN, self).__init__()
+        self.memory_dim = memory_dim
+        self.edge_feat_dim = edge_feat_dim
+        self.temporal_dim = temporal_dim
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.n_neighbors = n_neighbors
+        self.memory_updater_type = memory_updater_type
+        self.num_nodes = num_nodes
+
+        self.memory = MemoryModule(self.num_nodes,
+                                   self.memory_dim)
+
+        self.memory_ops = MemoryOperation(self.memory_updater_type,
+                                          self.memory,
+                                          self.edge_feat_dim,
+                                          self.temporal_dim)
+
+        
+        self.embedding_attn = TemporalGATConv(self.edge_feat_dim,
+                                              self.memory_dim,
+                                              self.temporal_dim,
+                                              self.embedding_dim,
+                                              self.num_heads,
+                                              allow_zero_in_degree=True)
+
+        self.msg_linkpredictor = MsgLinkPredictor(embedding_dim)
+
+    def embed(self, postive_graph, negative_graph, blocks):
+        emb_graph = blocks[0]
+        emb_memory = self.memory.memory[emb_graph.ndata[dgl.NID], :]
+        emb_t = emb_graph.ndata['timestamp']
+        embedding = self.embedding_attn(emb_graph, emb_memory, emb_t)
+        emb2pred = dict(
+            zip(emb_graph.ndata[dgl.NID].tolist(), emb_graph.nodes().tolist()))
+        # Since postive graph and negative graph has same is mapping
+        feat_id = [emb2pred[int(n)] for n in postive_graph.ndata[dgl.NID]]
+        feat = embedding[feat_id]
+        pred_pos, pred_neg = self.msg_linkpredictor(
+            feat, postive_graph, negative_graph)
+        return pred_pos, pred_neg
+
+    def update_memory(self, subg):
+        new_g = self.memory_ops(subg)
+        self.memory.set_memory(new_g.ndata[dgl.NID], new_g.ndata['memory'])
+        self.memory.set_last_update_t(
+            new_g.ndata[dgl.NID], new_g.ndata['timestamp'])
+
+    # Some memory operation wrappers
+    def detach_memory(self):
+        self.memory.detach_memory()
+
+    def reset_memory(self):
+        self.memory.reset_memory()
+
+    def store_memory(self):
+        memory_checkpoint = {}
+        memory_checkpoint['memory'] = copy.deepcopy(self.memory.memory)
+        memory_checkpoint['last_t'] = copy.deepcopy(self.memory.last_update_t)
+        return memory_checkpoint
+
+    def restore_memory(self, memory_checkpoint):
+        self.memory.memory = memory_checkpoint['memory']
+        self.memory.last_update_time = memory_checkpoint['last_t']

examples/pytorch/tgn/train.py

+import argparse
+import traceback
+import time
+import copy
+
+import numpy as np
+import dgl
+import torch
+
+from tgn import TGN
+from data_preprocess import TemporalWikipediaDataset, TemporalRedditDataset, TemporalDataset
+from dataloading import (FastTemporalEdgeCollator, FastTemporalSampler,
+                         SimpleTemporalEdgeCollator, SimpleTemporalSampler,
+                         TemporalEdgeDataLoader, TemporalSampler, TemporalEdgeCollator)
+
+from sklearn.metrics import average_precision_score, roc_auc_score
+
+
+TRAIN_SPLIT = 0.7
+VALID_SPLIT = 0.85
+
+# set random Seed
+np.random.seed(2021)
+torch.manual_seed(2021)
+
+
+def train(model, dataloader, sampler, criterion, optimizer, batch_size, fast_mode):
+    model.train()
+    total_loss = 0
+    batch_cnt = 0
+    last_t = time.time()
+    for _, positive_pair_g, negative_pair_g, blocks in dataloader:
+        optimizer.zero_grad()
+        pred_pos, pred_neg = model.embed(
+            positive_pair_g, negative_pair_g, blocks)
+        loss = criterion(pred_pos, torch.ones_like(pred_pos))
+        loss += criterion(pred_neg, torch.zeros_like(pred_neg))
+        total_loss += float(loss)*batch_size
+        retain_graph = True if batch_cnt == 0 and not fast_mode else False
+        loss.backward(retain_graph=retain_graph)
+        optimizer.step()
+        model.detach_memory()
+        model.update_memory(positive_pair_g)
+        if fast_mode:
+            sampler.attach_last_update(model.memory.last_update_t)
+        print("Batch: ", batch_cnt, "Time: ", time.time()-last_t)
+        last_t = time.time()
+        batch_cnt += 1
+    return total_loss
+
+
+def test_val(model, dataloader, sampler, criterion, batch_size, fast_mode):
+    model.eval()
+    batch_size = batch_size
+    total_loss = 0
+    aps, aucs = [], []
+    batch_cnt = 0
+    with torch.no_grad():
+        for _, postive_pair_g, negative_pair_g, blocks in dataloader:
+            pred_pos, pred_neg = model.embed(
+                postive_pair_g, negative_pair_g, blocks)
+            loss = criterion(pred_pos, torch.ones_like(pred_pos))
+            loss += criterion(pred_neg, torch.zeros_like(pred_neg))
+            total_loss += float(loss)*batch_size
+            y_pred = torch.cat([pred_pos, pred_neg], dim=0).sigmoid().cpu()
+            y_true = torch.cat(
+                [torch.ones(pred_pos.size(0)), torch.zeros(pred_neg.size(0))], dim=0)
+            model.update_memory(postive_pair_g)
+            if fast_mode:
+                sampler.attach_last_update(model.memory.last_update_t)
+            aps.append(average_precision_score(y_true, y_pred))
+            aucs.append(roc_auc_score(y_true, y_pred))
+            batch_cnt += 1
+    return float(torch.tensor(aps).mean()), float(torch.tensor(aucs).mean())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--epochs", type=int, default=50,
+                        help='epochs for training on entire dataset')
+    parser.add_argument("--batch_size", type=int,
+                        default=200, help="Size of each batch")
+    parser.add_argument("--embedding_dim", type=int, default=100,
+                        help="Embedding dim for link prediction")
+    parser.add_argument("--memory_dim", type=int, default=100,
+                        help="dimension of memory")
+    parser.add_argument("--temporal_dim", type=int, default=100,
+                        help="Temporal dimension for time encoding")
+    parser.add_argument("--memory_updater", type=str, default='gru',
+                        help="Recurrent unit for memory update")
+    parser.add_argument("--aggregator", type=str, default='last',
+                        help="Aggregation method for memory update")
+    parser.add_argument("--n_neighbors", type=int, default=10,
+                        help="number of neighbors while doing embedding")
+    parser.add_argument("--sampling_method", type=str, default='topk',
+                        help="In embedding how node aggregate from its neighor")
+    parser.add_argument("--num_heads", type=int, default=8,
+                        help="Number of heads for multihead attention mechanism")
+    parser.add_argument("--fast_mode", action="store_true", default=False,
+                        help="Fast Mode uses batch temporal sampling, history within same batch cannot be obtained")
+    parser.add_argument("--simple_mode", action="store_true", default=False,
+                        help="Simple Mode directly delete the temporal edges from the original static graph")
+    parser.add_argument("--num_negative_samples", type=int, default=1,
+                        help="number of negative samplers per positive samples")
+    parser.add_argument("--dataset", type=str, default="wikipedia",
+                        help="dataset selection wikipedia/reddit")
+    parser.add_argument("--k_hop", type=int, default=1,
+                        help="sampling k-hop neighborhood")                        
+
+    args = parser.parse_args()
+
+    assert not (args.fast_mode and args.simple_mode), "you can only choose one sampling mode"
+    if args.k_hop != 1:
+        assert args.simple_mode, "this k-hop parameter only support simple mode"
+
+    if args.dataset == 'wikipedia':
+        data = TemporalWikipediaDataset()
+    elif args.dataset == 'reddit':
+        data = TemporalRedditDataset()
+    else:
+        print("Warning Using Untested Dataset: "+args.dataset)
+        data = TemporalDataset(args.dataset)
+
+    # Pre-process data, mask new node in test set from original graph
+    num_nodes = data.num_nodes()
+    num_edges = data.num_edges()
+
+    num_edges = data.num_edges()
+    trainval_div = int(VALID_SPLIT*num_edges)
+
+    # Select new node from test set and remove them from entire graph
+    test_split_ts = data.edata['timestamp'][trainval_div]
+    test_nodes = torch.cat([data.edges()[0][trainval_div:], data.edges()[
+                           1][trainval_div:]]).unique().numpy()
+    test_new_nodes = np.random.choice(
+        test_nodes, int(0.1*len(test_nodes)), replace=False)
+
+    in_subg = dgl.in_subgraph(data, test_new_nodes)
+    out_subg = dgl.out_subgraph(data, test_new_nodes)
+    # Remove edge who happen before the test set to prevent from learning the connection info
+    new_node_in_eid_delete = in_subg.edata[dgl.EID][in_subg.edata['timestamp'] < test_split_ts]
+    new_node_out_eid_delete = out_subg.edata[dgl.EID][out_subg.edata['timestamp'] < test_split_ts]
+    new_node_eid_delete = torch.cat(
+        [new_node_in_eid_delete, new_node_out_eid_delete]).unique()
+
+    graph_new_node = copy.deepcopy(data)
+    # relative order preseved
+    graph_new_node.remove_edges(new_node_eid_delete)
+
+    # Now for no new node graph, all edge id need to be removed
+    in_eid_delete = in_subg.edata[dgl.EID]
+    out_eid_delete = out_subg.edata[dgl.EID]
+    eid_delete = torch.cat([in_eid_delete, out_eid_delete]).unique()
+
+    graph_no_new_node = copy.deepcopy(data)
+    graph_no_new_node.remove_edges(eid_delete)
+
+    # graph_no_new_node and graph_new_node should have same set of nid
+
+    # Sampler Initialization
+    if args.simple_mode:
+        fan_out = [args.n_neighbors for _ in range(args.k_hop)]
+        sampler = SimpleTemporalSampler(graph_no_new_node, fan_out)
+        new_node_sampler = SimpleTemporalSampler(data, fan_out)
+        edge_collator = SimpleTemporalEdgeCollator
+    elif args.fast_mode:
+        sampler = FastTemporalSampler(graph_no_new_node, k=args.n_neighbors)
+        new_node_sampler = FastTemporalSampler(data, k=args.n_neighbors)
+        edge_collator = FastTemporalEdgeCollator
+    else:
+        sampler = TemporalSampler(k=args.n_neighbors)
+        edge_collator = TemporalEdgeCollator
+
+    neg_sampler = dgl.dataloading.negative_sampler.Uniform(
+        k=args.num_negative_samples)
+    # Set Train, validation, test and new node test id
+    train_seed = torch.arange(int(TRAIN_SPLIT*graph_no_new_node.num_edges()))
+    valid_seed = torch.arange(int(
+        TRAIN_SPLIT*graph_no_new_node.num_edges()), trainval_div-new_node_eid_delete.size(0))
+    test_seed = torch.arange(
+        trainval_div-new_node_eid_delete.size(0), graph_no_new_node.num_edges())
+    test_new_node_seed = torch.arange(
+        trainval_div-new_node_eid_delete.size(0), graph_new_node.num_edges())
+
+    g_sampling = None if args.fast_mode else dgl.add_reverse_edges(
+        graph_no_new_node, copy_edata=True)
+    new_node_g_sampling = None if args.fast_mode else dgl.add_reverse_edges(
+        graph_new_node, copy_edata=True)
+    if not args.fast_mode:
+        new_node_g_sampling.ndata[dgl.NID] = new_node_g_sampling.nodes()
+        g_sampling.ndata[dgl.NID] = new_node_g_sampling.nodes()
+
+    # we highly recommend that you always set the num_workers=0, otherwise the sampled subgraph may not be correct.
+    train_dataloader = TemporalEdgeDataLoader(graph_no_new_node,
+                                              train_seed,
+                                              sampler,
+                                              batch_size=args.batch_size,
+                                              negative_sampler=neg_sampler,
+                                              shuffle=False,
+                                              drop_last=False,
+                                              num_workers=0,
+                                              collator=edge_collator,
+                                              g_sampling=g_sampling)
+
+    valid_dataloader = TemporalEdgeDataLoader(graph_no_new_node,
+                                              valid_seed,
+                                              sampler,
+                                              batch_size=args.batch_size,
+                                              negative_sampler=neg_sampler,
+                                              shuffle=False,
+                                              drop_last=False,
+                                              num_workers=0,
+                                              collator=edge_collator,
+                                              g_sampling=g_sampling)
+
+    test_dataloader = TemporalEdgeDataLoader(graph_no_new_node,
+                                             test_seed,
+                                             sampler,
+                                             batch_size=args.batch_size,
+                                             negative_sampler=neg_sampler,
+                                             shuffle=False,
+                                             drop_last=False,
+                                             num_workers=0,
+                                             collator=edge_collator,
+                                             g_sampling=g_sampling)
+
+    test_new_node_dataloader = TemporalEdgeDataLoader(graph_new_node,
+                                                      test_new_node_seed,
+                                                      new_node_sampler if args.fast_mode else sampler,
+                                                      batch_size=args.batch_size,
+                                                      negative_sampler=neg_sampler,
+                                                      shuffle=False,
+                                                      drop_last=False,
+                                                      num_workers=0,
+                                                      collator=edge_collator,
+                                                      g_sampling=new_node_g_sampling)
+
+    edge_dim = data.edata['feats'].shape[1]
+    num_node = data.num_nodes()
+
+    model = TGN(edge_feat_dim=edge_dim,
+                memory_dim=args.memory_dim,
+                temporal_dim=args.temporal_dim,
+                embedding_dim=args.embedding_dim,
+                num_heads=args.num_heads,
+                num_nodes=num_node,
+                n_neighbors=args.n_neighbors,
+                memory_updater_type=args.memory_updater)
+
+    criterion = torch.nn.BCEWithLogitsLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
+    # Implement Logging mechanism
+    f = open("logging.txt", 'w')
+    if args.fast_mode:
+        sampler.reset()
+    try:
+        for i in range(args.epochs):
+            train_loss = train(model, train_dataloader, sampler,
+                                criterion, optimizer, args.batch_size, args.fast_mode)
+            val_ap, val_auc = test_val(
+                model, valid_dataloader, sampler, criterion, args.batch_size, args.fast_mode)
+            memory_checkpoint = model.store_memory()
+            if args.fast_mode:
+                new_node_sampler.sync(sampler)
+            test_ap, test_auc = test_val(
+                model, test_dataloader, sampler, criterion, args.batch_size, args.fast_mode)
+            model.restore_memory(memory_checkpoint)
+            if args.fast_mode:
+                sample_nn = new_node_sampler
+            else:
+                sample_nn = sampler
+            nn_test_ap, nn_test_auc = test_val(
+                model, test_new_node_dataloader, sample_nn, criterion, args.batch_size, args.fast_mode)
+            log_content = []
+            log_content.append("Epoch: {}; Training Loss: {} | Validation AP: {:.3f} AUC: {:.3f}\n".format(
+                i, train_loss, val_ap, val_auc))
+            log_content.append(
+                "Epoch: {}; Test AP: {:.3f} AUC: {:.3f}\n".format(i, test_ap, test_auc))
+            log_content.append("Epoch: {}; Test New Node AP: {:.3f} AUC: {:.3f}\n".format(
+                i, nn_test_ap, nn_test_auc))
+
+            f.writelines(log_content)
+            model.reset_memory()
+            if i < args.epochs-1 and args.fast_mode:
+                sampler.reset()
+            print(log_content[0], log_content[1], log_content[2])
+    except:
+        traceback.print_exc()
+        error_content = "Training Interreputed!"
+        f.writelines(error_content)
+        f.close()
+        # exit(-1)
+    print("========Training is Done========")