Skip to content

PyTorch implementation of IndRNN with custom C++/CUDA extensions

Notifications You must be signed in to change notification settings

ysngshn/torch-indrnn

Repository files navigation

This is a PyTorch implementation of the IndRNN architecture with custom C++/CUDA extensions for efficiency.

IndRNN is a simple and effective recurrent neural network structure proposed by Li et. al. at CVPR 2018:

"Independently Recurrent Neural Network (IndRNN): Building a Longer and Deeper RNN". Shuai Li, Wanqing Li, Chris Cook, Ce Zhu, Yanbo Gao. CVPR 2018

The authors have also provided an official PyTorch implementation:

github.com/Sunnydreamrain/IndRNN_pytorch

Compared to the official implementation, this package is a clean implementation that follows the standard PyTorch C++/CUDA API. This should result in better compatibility with PyTorch, supporting for instance computation at different precision. Also, this implementation has no dependency on CuPy or pynvrtc.

Installation

To use this package, make sure you have installed PyTorch and have a compatible cuda version with nvcc compiler.

Then you can install this library locally via

pip install -e .