[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/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========")