PredRNN：用于时空预测学习的循环神经网络（TPAMI 2022）

时空序列的预测学习旨在通过学习历史上下文来生成未来的图像，其中视觉动态被认为具有模块化的结构，可以通过组合式的子系统来学习。

初版于 NeurIPS 2017

本仓库首先包含 PredRNN（2017）的 PyTorch 实现 [论文]，这是一种具有两组记忆单元的循环网络，它们以几乎独立的方式进行状态转移，最终形成对复杂环境的统一表征。

具体而言，除了 LSTM 的原始记忆单元外，该网络还具有一个在所有层中自下而上和自上而下双向传播的锯齿形记忆流，从而使不同层次 RNN 学习到的视觉动态能够相互交流。

TPAMI 2022 版 PredRNN-V2 的新特性

本仓库还包括 PredRNN-V2 的实现 [论文]，它在以下三个方面对 PredRNN 进行了改进。

1. 记忆解耦的 ST-LSTM

我们发现，PredRNN 中的两组记忆单元包含不必要的冗余特征，因此提出了一种记忆解耦损失，以鼓励它们学习视觉动态的模块化结构。

2. 反向调度采样

反向调度采样是一种针对序列到序列 RNN 的新型课程学习策略。与传统的调度采样相反，它逐渐将 PredRNN 编码器的训练过程从使用先前生成的帧切换为使用之前的真值帧。优点：这迫使模型从上下文帧中学习长期动态。

3. 动作条件下的视频预测

我们进一步将 PredRNN 扩展到动作条件下的视频预测任务。通过将动作与隐藏状态融合，PredRNN 和 PredRNN-V2 在长期预测方面表现出极高的竞争力，有望作为基于模型的视觉控制中的基础动力学模型。

我们在 BAIR 机器人推动物体数据集上展示了定量结果，该任务是从 2 个观测帧预测未来 28 帧。

展示案例

移动 MNIST

KTH

BAIR（我们放大红色框内的区域）

Traffic4Cast

雷达回波

LPIPS 下移动 MNIST 和 KTH 的定量结果

LPIPS 对人类感知判断更为敏感，数值越低越好。

	移动 MNIST	KTH 动作
PredRNN	0.109	0.204
PredRNN-V2	0.071	0.139

Traffic4Cast（柏林）的定量结果

	MSE (10^{-3})
U-Net	6.992
CrevNet	6.789
U-Net+PredRNN-V2	5.135

快速入门

1. 安装 Python 3.6 和 PyTorch 1.9.0 以运行主代码。同时，安装 Tensorflow 2.1.0 用于 BAIR 数据加载器。

2. 下载数据。本仓库包含三个数据集的代码：移动 MNIST 数据集、KTH 动作数据集，以及 BAIR 数据集（30.1GB），可通过以下命令获取：

   wget http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar
   

3. 训练模型。您可以使用以下 Bash 脚本来训练模型。训练好的模型将保存在 --save_dir 文件夹中，生成的未来帧将保存在 --gen_frm_dir 文件夹中。

4. 您可以从 清华大学云 或 Google Drive 获取 预训练模型。

cd mnist_script/
sh predrnn_mnist_train.sh
sh predrnn_v2_mnist_train.sh

cd kth_script/
sh predrnn_kth_train.sh
sh predrnn_v2_kth_train.sh

cd bair_script/
sh predrnn_bair_train.sh
sh predrnn_v2_bair_train.sh

引用

如果您觉得本仓库有用，请引用以下论文。

@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 Yu, Philip S},
  booktitle={Advances in Neural Information Processing Systems},
  pages={879--888},
  year={2017}
}

@misc{wang2021predrnn,
      title={{PredRNN}: A Recurrent Neural Network for Spatiotemporal Predictive Learning}, 
      author={Wang, Yunbo and Wu, Haixu and Zhang, Jianjin and Gao, Zhifeng and Wang, Jianmin and Yu, Philip S and Long, Mingsheng},
      year={2021},
      eprint={2103.09504},
      archivePrefix={arXiv},
}

PredRNN-PyTorch 快速上手指南

PredRNN 是一个用于时空序列预测学习的循环神经网络，适用于视频预测、雷达回波外推等任务。本指南基于官方仓库，帮助您快速在本地运行 PredRNN 及改进版 PredRNN-V2。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python 版本: 3.6
• 核心依赖:
  • PyTorch 1.9.0
  • TensorFlow 2.1.0 (仅用于 BAIR 数据集的数据加载器)
• 硬件建议: 建议使用支持 CUDA 的 NVIDIA GPU 以加速训练和推理。

安装步骤

1. 创建虚拟环境并安装 Python 依赖 建议使用 conda 或 venv 隔离环境。以下以 conda 为例：

   conda create -n predrnn python=3.6
   conda activate predrnn
   pip install torch==1.9.0 torchvision==0.10.0
   pip install tensorflow==2.1.0
   pip install numpy scipy scikit-image opencv-python
   

   提示：国内用户可使用清华源加速安装： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package_name>

2. 克隆项目代码

   git clone https://github.com/violanteo/predrnn-pytorch.git
   cd predrnn-pytorch
   

3. 下载数据集 根据需求选择以下任一数据集：

   • Moving MNIST: 下载链接
   • KTH Action: 下载链接
   • BAIR Robot Pushing (30.1GB):

     wget http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar
     tar -xvf bair_robot_pushing_dataset_v0.tar
     

   请将下载的数据集解压后放置在代码中指定的数据目录（通常在脚本中通过参数配置）。

4. 获取预训练模型（可选） 若需直接测试或微调，可从以下地址下载预训练权重：

   • 清华大学云盘 (国内推荐)
   • Google Drive

基本使用

项目提供了针对不同数据集的训练脚本。以下以 Moving MNIST 数据集为例，演示如何训练原始版 PredRNN 和改进版 PredRNN-V2。

1. 训练 Moving MNIST 模型

进入对应的脚本目录并执行训练命令：

cd mnist_script/

# 训练原始版 PredRNN
sh predrnn_mnist_train.sh

# 训练改进版 PredRNN-V2
sh predrnn_v2_mnist_train.sh

• 训练好的模型将保存在脚本参数 --save_dir 指定的文件夹中。
• 生成的未来预测帧将保存在 --gen_frm_dir 指定的文件夹中。

2. 训练其他数据集

对于 KTH 动作数据集或 BAIR 机器人推动数据集，操作方式类似：

# KTH 数据集
cd ../kth_script/
sh predrnn_kth_train.sh
sh predrnn_v2_kth_train.sh

# BAIR 数据集
cd ../bair_script/
sh predrnn_bair_train.sh
sh predrnn_v2_bair_train.sh

3. 结果评估

训练完成后，您可以查看生成目录中的图像序列，或参考 README 中提供的 LPIPS/MSE 指标对比表来评估模型性能。PredRNN-V2 通常在长时预测任务中表现更优。