{"id":13936131,"url":"https://github.com/L0SG/relational-rnn-pytorch","last_synced_at":"2025-07-19T21:31:49.403Z","repository":{"id":111962629,"uuid":"145528100","full_name":"L0SG/relational-rnn-pytorch","owner":"L0SG","description":"An implementation of DeepMind's Relational Recurrent Neural Networks (NeurIPS 2018) in PyTorch.","archived":false,"fork":false,"pushed_at":"2018-12-27T05:38:23.000Z","size":4707,"stargazers_count":245,"open_issues_count":1,"forks_count":35,"subscribers_count":10,"default_branch":"master","last_synced_at":"2025-07-11T07:59:37.558Z","etag":null,"topics":["deep-learning","deepmind","language-model","language-modeling","pytorch","recurrent-neural-networks","self-attention","transformer","word-language-model"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/L0SG.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2018-08-21T07:57:41.000Z","updated_at":"2024-11-19T02:51:44.000Z","dependencies_parsed_at":"2023-06-17T15:33:52.338Z","dependency_job_id":null,"html_url":"https://github.com/L0SG/relational-rnn-pytorch","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/L0SG/relational-rnn-pytorch","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/L0SG%2Frelational-rnn-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/L0SG%2Frelational-rnn-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/L0SG%2Frelational-rnn-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/L0SG%2Frelational-rnn-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/L0SG","download_url":"https://codeload.github.com/L0SG/relational-rnn-pytorch/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/L0SG%2Frelational-rnn-pytorch/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":266019657,"owners_count":23864916,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","deepmind","language-model","language-modeling","pytorch","recurrent-neural-networks","self-attention","transformer","word-language-model"],"created_at":"2024-08-07T23:02:24.450Z","updated_at":"2025-07-19T21:31:44.386Z","avatar_url":"https://github.com/L0SG.png","language":"Python","funding_links":[],"categories":["Python","Paper implementations｜论文实现","Paper implementations"],"sub_categories":["Other libraries｜其他库:","Other libraries:"],"readme":"# relational-rnn-pytorch\n\nAn implementation of DeepMind's [Relational Recurrent Neural Networks](https://arxiv.org/abs/1806.01822) (Santoro et al. 2018) in PyTorch.\n\n![](./pics/rmc.png)\n![](./pics/rmc_paper_result.png)\n\n\nRelational Memory Core (RMC) module is originally from [official Sonnet implementation](https://github.com/deepmind/sonnet/blob/master/sonnet/python/modules/relational_memory.py). However, currently they do not provide a full language modeling benchmark code.\n\nThis repo is a port of RMC with additional comments. It features a full-fledged word language modeling benchmark vs. traditional LSTM.\n\nIt supports any arbitrary word token-based text dataset, including WikiText-2 \u0026 WikiText-103.\n\nBoth RMC \u0026 LSTM models support [adaptive softmax](https://pytorch.org/docs/stable/nn.html#adaptivelogsoftmaxwithloss) for much lower memory usage of large vocabulary dataset. RMC supports PyTorch's `DataParallel`, so you can easily experiment with a multi-GPU setup.\n\nbenchmark codes are hard-forked from [official PyTorch word-language-model example](https://github.com/pytorch/examples/tree/master/word_language_model)\n\nIt also features an N-th farthest synthetic task from the paper (see below).\n\n# Requirements\nPyTorch 0.4.1 or later (Tested on 1.0.0) \u0026 Python 3.6\n\n# Examples\n`python train_rmc.py --cuda ` for full training \u0026 test run of RMC with GPU.\n\n`python train_rmc.py --cuda --adaptivesoftmax --cutoffs 1000 5000 20000` if using large vocabulary dataset (like WikiText-103) to fit all the tensors in the VRAM.\n\n`python generate_rmc.py --cuda` for generating sentences from the trained model.\n\n`python train_rnn.py --cuda` for full training \u0026 test run of traditional RNN with GPU.\n\nAll default hyperparameters of RMC \u0026 LSTM are results from a two-week experiment using WikiText-2.\n\n# Data Preparation\nTested with WikiText-2 and WikiText-103. WikiText-2 is bundled.\n\nCreate a subfolder inside `./data` and place word-level `train.txt`, `valid.txt`, and `test.txt` inside the subfolder.\n\nSpecify `--data=(subfolder name)` and you are good to go.\n\nThe code performs tokenization at the first training run, and the corpus is saved as `pickle`. The code will load the `pickle` file after the first run.\n\n# WikiText-2 Benchmark Results\nBoth RMC \u0026 LSTM have ~11M parameters. Please refer to the training code for details on hyperparameters.\n\n| Models        | Valid Perplexity|Test Perplexity           | Forward pass ms/batch (TITAN Xp) |  Forward pass ms/batch (TITAN V) |\n|:-------------:|:-------------:|:-------------:| :-------------:| :-------------:|\n| LSTM (CuDNN)      |111.31 | 105.56 | 26~27 | 40~41 |\n| LSTM (For Loop)      |Same as CuDNN | Same as CuDNN | 30~31 | 60~61 |\n| RMC      | 112.77 | 107.21      |  110~130  | 220~230|\n\nRMC can reach a comparable performance to LSTM (with heavy hyperparameter search), but it turns out that the RMC is very slow. The multi-head self-attention at every time step may be the culprit here.\nUsing LSTMCell with for loop (which is more \"fair\" benchmark for RMC) slows down the forward pass, but it's still much faster.\n\nPlease also note that the hyperparameter for RMC is a worst-case scenario in terms of speed, because it used a single memory slot (as described in the paper) and did not benefit from a row-wise weight sharing from multi-slot memory.  \n\nInteresting to note here is that the speed is slower in TITAN V than TITAN Xp. The reason might be that the models are relatively small and the model calls small linear operations frequently.\n\nMaybe TITAN Xp (~1,900Mhz unlocked CUDA clock speed vs. TITAN V's 1,335Mhz limit) benefits from these kind of workload. Or maybe TITAN V's CUDA kernel launch latency is higher for the ops in the model.\n\nI'm not an expert in details of CUDA. Please share your results!  \n\n# RMC Hyperparameter Search Results\nAttention parameters tend to overfit the WikiText-2. reducing the hyperparmeters for attention (key_size) can combat the overfitting.\n\nApplying dropout at the output logit before the softmax (like the LSTM one) helped preventing the overfitting.\n\n|embed \u0026 head size| # heads | attention MLP layers | key size | dropout at output | memory slots | test ppl|\n|:----:|:----:|:----:|:----:|:----:|:----:|:----:|\n|128|\t4|\t3|\t128|\tNo|\t1|\t128.81 |\n|128|\t4|\t3|\t128|\tNo|\t1|\t128.81 |\n|128|\t8|\t3|\t128|\tNo|\t1|\t141.84 |\n|128|\t4|\t3|\t32|\tNo\t|1\t|123.26 |\n|128|\t4|\t3|\t32|\tYes|\t1|\t112.4 |\n|128|\t4|\t3|\t64|\tNo\t|1\t|124.44 |\n|128|\t4|\t3|\t64|\tYes|\t1|\t110.16 |\n|128|\t4|\t2|\t64|\tYes|\t1|\t111.67 |\n|64\t|4\t|3\t|64\t|Yes\t|1\t|133.68 |\n|64\t|4\t|3\t|32\t|Yes\t|1\t|135.93 |\n|64\t|4\t|3\t|64\t|Yes\t|4\t|137.93 |\n|192|\t4|\t3|\t64|\tYes|\t1|\t**107.21** |\n|192|\t4|\t3|\t64|\tYes|\t4|\t114.85 |\n|256|\t4|\t3|\t256|\tNo|\t1|\t194.73 |\n|256|\t4|\t3|\t64|\tYes|\t1|\t126.39 |\n\n\n# About WikiText-103\nThe original RMC paper presents WikiText-103 results with a larger model \u0026 batch size (6 Tesla P100, each with 64 batch size, so a total of 384. Ouch).\n\nUsing a full softmax easily blows up the VRAM. Using `--adaptivesoftmax` is highly recommended. If using `--adaptivesoftmax`, `--cutoffs` should be properly provided. Please refer to the [original API description](https://pytorch.org/docs/stable/nn.html#adaptivelogsoftmaxwithloss)\n\nI don't have such hardware and my resource is too limited to do the experiments. Benchmark result, or any other contributions are very welcome!\n\n# Nth Farthest Task\n\nThe objective of the task is: Given k randomly labelled (from 1 to k) D-dimensional vectors, identify which is the Nth farthest vector from vector M. (The answer is an integer from 1 to k.)\n\nThe specific task in the paper is: given 8 labelled 16-dimensional vectors, which is the Nth farthest vector from vector M? The vectors are labelled randomly so the model has to recognise that the Mth vector is the vector labelled as M as opposed to the vector in the Mth position in the input. \n\nThe input to the model comprises 8 40-dimensional vectors for each example. Each of these 40-dimensional vectors is structured like this:\n\n```\n[(vector 1) (label: which vector is it, from 1 to 8, one-hot encoded) (N, one-hot encoded) (M, one-hot encoded)] \n```\n\n#### Example\n\n`python train_nth_farthest.py --cuda` for training and testing on the Nth Farthest Task with GPU(s).\n\nThis uses the `RelationalMemory` class in `relational_rnn_general.py`, which is a version of `relational_rnn_models.py` without the language-modelling specific code.\n\nPlease refer to`train_nth_farthest.py` for details on hyperparameter values. These are taken from Appendix A1 in the paper and from the Sonnet implementation when the hyperparameter values are not given in the paper.\n\nNote: new examples are generated per epoch as in the Sonnet implementation. This seems to be consistent with the paper, which does not specify the number of examples used.\n\n#### Experiment results\n\nThe model has been trained with a single TITAN Xp GPU for forever until it reaches 91% test accuracy. Below are the results with 3 independent runs:\n![](./pics/nth_results.jpg)\n\nThe model does break the 25% barrier if trained long enough, but the wall clock time is roughly over 2~3x longer than those reported in the paper.\n\n#### TODO\n\nExperiment with different hyperparameters \n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FL0SG%2Frelational-rnn-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2FL0SG%2Frelational-rnn-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FL0SG%2Frelational-rnn-pytorch/lists"}