Skip to main content

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

AWS AWS Cloud computing Datadog Datadog Monitoring DigiCert DigiCert EV certificate Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page