Skip to main content
Join the official 2020 Python Developers SurveyStart the survey!

Implements adaptive computation time RNNs in PyTorch, with the same interface as builtin RNNs.

Project description


This library implements PyTorch modules for recurrent neural networks that can learn to execute variable-time algorithms, as presented in Adaptive Computation Time for Recurrent Neural Networks (Graves 2016). These models can learn patterns requiring varying amounts of computation for a fixed-size input, which is difficult or impossible for traditional neural networks. The library aims to be clean, idiomatic, and extensible, offering a similar interface to PyTorch’s builtin recurrent modules.

The main features are:

  • A nearly drop-in replacement for torch.nn.RNN- and torch.nn.RNNCell-style RNNs, but with the power of variable computation time.
  • A wrapper which adds adaptive computation time to any RNNCell.
  • Data generators, configs, and training scripts to reproduce experiments from the paper.


Vanilla PyTorch GRU:

rnn = torch.nn.GRU(64, 128, num_layers=2)
output, hidden = rnn(inputs, initial_hidden)

GRU with adaptive computation time:

rnn = models.AdaptiveGRU(64, 128, num_layers=2, time_penalty=1e-3)
output, hidden, ponder_cost = rnn(inputs, initial_hidden)


Documentation is hosted on Read the Docs.


You don’t need to cite this code, but if it helps you in your research and you’d like to:

  title   = "pytorch-adaptive-computation-time",
  author  = "Swope, Aidan",
  journal = "GitHub",
  year    = "2020",
  url     = ""

If you use the experiment code, please also consider citing PyTorch Lightning.

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for pytorch-adaptive-computation-time, version 0.1.3
Filename, size File type Python version Upload date Hashes
Filename, size pytorch_adaptive_computation_time-0.1.3-py3-none-any.whl (15.7 kB) File type Wheel Python version py3 Upload date Hashes View
Filename, size pytorch-adaptive-computation-time-0.1.3.tar.gz (14.2 kB) File type Source Python version None Upload date Hashes View

Supported by

Pingdom Pingdom Monitoring Google Google Object Storage and Download Analytics Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN DigiCert DigiCert EV certificate StatusPage StatusPage Status page