zzwei1 / predrnn-pytorch

Official implementation for NIPS'17 paper: PredRNN: Recurrent Neural Networks for Predictive Learning Using Spatiotemporal LSTMs.

PredRNN (NIPS 2017)

A PyTorch implementation of PredRNN [paper], a recurrent network with twisted and zigzag space-time memory cells for video data. Given a sequence of previous frames, our model generates future frames for multiple timestamps.

Video prediction networks have been used for precipitation nowcasting, early activity recognition, physical scene understanding, model-based visual planning, and unsupervised representation learning of video data.

Get Started

1. Install Python 3.7, PyTorch 1.3, and OpenCV 3.4.

2. Download data. This repo contains code for two datasets: the Moving Mnist dataset and the KTH action dataset.

3. Train the model. You can use the following bash script to train the model. The learned model will be saved in the --save_dir folder. The generated future frames will be saved in the --gen_frm_dir folder.

cd script/
sh predrnn_mnist_train.sh

Citation

If you use this repo or our results in your research, please remember to cite the following paper.

@inproceedings{wang2017predrnn,
  title={{PredRNN}: Recurrent Neural Networks for Predictive Learning Using Spatiotemporal {LSTM}s},
  author={Wang, Yunbo and Long, Mingsheng and Wang, Jianmin and Gao, Zhifeng and Philip, S Yu},
  booktitle={Advances in Neural Information Processing Systems},
  pages={879--888},
  year={2017}
}

Related Publication

PredRNN++: Towards A Resolution of the Deep-in-Time Dilemma in Spatiotemporal Predictive Learning.
Yunbo Wang, Zhifeng Gao, Mingsheng Long, Jianmin Wang, and Philip S. Yu.
ICML 2018 [paper] [code]

Contact

You may send email to [email protected] or [email protected], or create an issue in this repo and @wyb15.