An implementation of the TGN in TensorFlow 1 for IPU.
The temporal graph network is a dynamic GNN that can predict link edit operations in dynamic graphs. The model includes a sparsely accessed memory and a Graph Transformer layer. Following the authors, we apply it to a dataset of Wikipedia edits from JODIE. The task is to predict edit events that connect users with pages. We apply it in the transductive setting where the same set of users and pages is used for validation and testing (although the edit events are distinct).
The JODIE-Wikipedia dataset is available from snap.stanford.edu, but note that it will be automatically downloaded when run_tgn.py is first run.
We reproduce the behaviour of PyTorch Geometric examples/tgn.py. While the main body of our implementation is distinct, we share code for data loading and nearest neighbour preprocessing. We have also made a few modifications:
- Use lower precision (
tf.float16) where possible. - Recompute the memory at validation/test time in the same way as training (note: this should not change the results).
- Concatenate the memory payload to reduce the number of
tf.gathercalls. - When benchmarking, cache the dataset (reusing negative samples) and only validate at the end of training.
1. Download and install the Poplar SDK following the Getting Started guide for your IPU system.
2. Create a virtual environment and install the appropriate Graphcore TensorFlow 1.15 wheel from inside the SDK directory:
virtualenv --python python3 .venv
source .venv/bin/activate
source <path to enable.sh for Poplar from the Poplar SDK>
pip install <path to the tensorflow-1 wheel from the Poplar SDK>
pip install -r requirements.txt3. Run training, profiling or benchmarking:
# Training (should reach 0.97 auc, 0.97 ap, 0.35 training loss)
# NOTE: On first run, this will download the JODIE-Wikipedia dataset (approx 500M) - this may take a while
python run_tgn.py
# Profiling
env POPLAR_ENGINE_OPTIONS='{"autoReport.all":"true", "autoReport.directory":"report"}' python run_tgn.py -m profile| Path | Description |
|---|---|
run_tgn.py |
training loop and entry point for running the TGN |
model.py |
Tensorflow model and single training step |
optimiser.py |
low-precision Adam optimiser |
utils.py |
helpers for defining portable training loops |
dataloader.py |
dataset download, batching and neighbour sampling |
requirements.txt, setup.cfg |
setup & development |
tests/ |
unit tests |
data/JODIE/ |
dataset, automatically downloaded when running run_tgn.py |
- Temporal Graph Networks for Deep Learning on Dynamic Graphs, Rossi, E., et al., https://arxiv.org/abs/2006.10637
- Fast Graph Representation Learning with PyTorch Geometric, Fey, M. and Lenssen, J.E., https://arxiv.org/abs/1903.02428
- Predicting Dynamic Embedding Trajectory in Temporal Interaction Networks, Kumar, S., Zhang, X. and Leskovec, J., https://arxiv.org/abs/1908.01207
- Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification, Shi, Y., et al., https://arxiv.org/abs/2009.03509
This example is licensed under the MIT license - see the LICENSE file at the top-level of this repository.
This directory includes derived work from the following:
PyTorch Geometric, https://github.com/pyg-team/pytorch_geometric/
Copyright (c) 2021 Matthias Fey, Jiaxuan You <matthias.fey@tu-dortmund.de, jiaxuan@cs.stanford.edu>
Licensed under the MIT License