evnn-pytorch 0.1.0

EVNN: a torch extension for custom event based RNN models.

EvNN: Event-based Neural Networks

EvNN is a CUDA and C++ implementation of event-based RNN layers with built-in DropConnect and Zoneout regularization. These layers are exposed through C++ and Pytorch APIs for easy integration into your own projects or machine learning frameworks. The code framework and base layers are adopted from Haste Library.

EGRU: Event-based Gated Recurrent Unit

Event based GRU was publised as a conference paper at ICLR 2023: Efficient recurrent architectures through activity sparsity and sparse back-propagation through time (notable-top-25%)

Illustration of EGRU. A. A single unit of the original GRU model adapted from Cho et al.. B: EGRU unit with event generating mechanism. C: Heaviside function and surrogate gradient. D: Forward state dynamics for two EGRU units ($i$ and $j$ ). E: Activity-sparse backward dynamics for two EGRU units ($i$ and $j$ ). (Note that we only have to backpropagate through units that were active or whose state was close to the threshold at each time step.)

Which RNN types are currently supported?

• GRU

What's included in this project?

• a PyTorch API (evnn_pytorch) for event based neural networks

Install

Here's what you'll need to get started:

• a CUDA Compute Capability 3.7+ GPU (required only if using GPU)
• CUDA Toolkit 10.0+ (required only if using GPU)
• PyTorch 1.3+ for PyTorch integration (GPU optional)
• BLAS or any BLAS-like library for CPU computation.
• Eigen 3 to build the C++ examples (optional)

Once you have the prerequisites, you can install with pip or by building the source code.

Building from source

Note

Currenty supported only on Linux, use Docker for building on Windows.

make evnn_pytorch # Build PyTorch API

If you built the PyTorch API, install it with pip:

pip install evnn_pytorch-*.whl

If the CUDA Toolkit that you're building against is not in /usr/local/cuda, you must specify the $CUDA_HOME environment variable before running make:

CUDA_HOME=/usr/local/cuda-10.2 make

Performance

Code for the experiments and benchmarks presented in the paper are published in benchmarks directory. Note that these benchmarks have additional dependencies as documented in benchmarks/requirements.txt

Documentation

PyTorch API

import torch
import evnn_pytorch as evnn

# setting use_custom_cuda=False makes the model use pytorch code instead of EvNN extension
egru_layer =  evnn.EGRU(input_size, hidden_size, zoneout=0.0, batch_first=True,
                        use_custom_cuda=True)

egru_layer.cuda()

# `x` is a CUDA tensor with shape [N,T,C]
x = torch.rand([5, 25, 128]).cuda()

y, state = egru_layer(x)

The PyTorch API is documented in docs/pytorch/evnn_pytorch.md.

Code layout

• docs/pytorch/: API reference documentation for evnn_pytorch
• frameworks/pytorch/: PyTorch API and custom op code
• lib/: CUDA kernels and C++ API
• validation/: scripts and tests to validate output and gradients of RNN layers
• benchmarks/: Experiments from ICLR 2023 Paper.

Testing

use python unittest with this command

• Numpy is required for testing

python -m unittest discover -p '*_test.py' -s validation

Note

Tests will fail if you set the the dimensions (batch_size,time_steps,input_size, hidden_size) too high, this is because floating point errors can accumulate and cause the units to generate events one timestep off. This causes the numerical tests to fail but the Neural Network training will work without any issues.

Implementation notes

• the EGRU is based on Haste GRU which is in turn based on 1406.1078v1 (same as cuDNN) rather than 1406.1078v3

References

1. Nanavati, Sharvil, ‘Haste: A Fast, Simple, and Open RNN Library’, 2020 https://github.com/lmnt-com/haste/

2. K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN encoder–decoder for statistical machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 1724–1734, Doha, Qatar, Oct. 2014. Association for Computational Linguistics. doi: 10.3115/v1/D14-1179. URL

Citing this work

To cite this work, please use the following BibTeX entry:

@inproceedings{
evnn2023,
title={Efficient recurrent architectures through activity sparsity and sparse back-propagation through time},
author={Anand Subramoney and Khaleelulla Khan Nazeer and Mark Sch{\"o}ne and Christian Mayr and David Kappel},
booktitle={The Eleventh International Conference on Learning Representations },
year={2023},
url={https://openreview.net/forum?id=lJdOlWg8td},
howpublished={https://github.com/KhaleelKhan/EvNN/}
}

License

Apache 2.0