RainMamba 视频去雨复现实录：从 AutoDL 开机到 PSNR 出分

[0. 复现一个深度学习项目，本质上是配齐三件套](#0. 复现一个深度学习项目，本质上是配齐三件套)
[1. AutoDL 环境准备](#1. AutoDL 环境准备)
- [1.1 选机与计费模式](#1.1 选机与计费模式)
- [1.2 环境配置命令](#1.2 环境配置命令)
- [1.3 编译 mamba / causal-conv1d（有卡模式下）](#1.3 编译 mamba / causal-conv1d（有卡模式下）)
- [1.4 依赖版本踩坑记录](#1.4 依赖版本踩坑记录)
- [1.5 环境验证](#1.5 环境验证)
[2. 预训练权重 VRDS.pth 下载](#2. 预训练权重 VRDS.pth 下载)
[3. 测试数据集下载与目录改造](#3. 测试数据集下载与目录改造)
- [3.1 下载](#3.1 下载)
- [3.2 目录结构改造](#3.2 目录结构改造)
[4. 推理：命令与代码](#4. 推理：命令与代码)
- [4.1 一条命令跑测试](#4.1 一条命令跑测试)
- [4.2 这条命令背后发生了什么](#4.2 这条命令背后发生了什么)
- [4.3 结果查看与复现判据](#4.3 结果查看与复现判据)
- [4.4 demo：对任意雨天视频去雨](#4.4 demo：对任意雨天视频去雨)
[5. 成本复盘](#5. 成本复盘)
[6. 小结](#6. 小结)

RainMamba 视频去雨复现实录：从 AutoDL 开机到 PSNR 出分

论文：RainMamba: Enhanced Locality Learning with State Space Models for Video Deraining（ACM MM'24 Oral）

代码：https://github.com/TonyHongtaoWu/RainMamba

目标：只复现推理（用作者预训练权重跑出测试集指标），不从头训练。

0. 复现一个深度学习项目，本质上是配齐三件套

深度学习项目 = 代码 + 数据集 + 权重。三者各有去处：

组件	内容	在哪里获取
代码	网络结构、训练/测试脚本、Mamba 算子	GitHub（`git clone`，仓库仅 41 MB）
权重	训练固化下来的参数（`.pth`，约 450 MB）	作者 Google Drive
数据集	测试集 lq（雨帧）/ gt（干净帧）配对	作者 OneDrive（HKUST-GZ）

为什么权重不放在 GitHub 里？单文件 100 MB 硬限制，模型权重动辄数百 MB，Git 仓库装不下。所以复现推理 = 克隆代码 + 下载权重 + 摆好数据目录 + 一条测试命令，四步缺一不可：只有代码没权重，函数结构已知但参数随机；只有权重没代码，参数是死数；没有数据集，跑得出图像但算不出 PSNR（评估需要 gt 作标准答案）。

1. AutoDL 环境准备

1.1 选机与计费模式

实操中总结的几条教训：

多卡实例按整机计费。租了 6×3090（¥9.36/时），跑推理只用 1 张卡，另外 5 张白烧钱。单卡任务应租单卡实例。
无卡模式（约 ¥0.1/时）适合装环境：克隆仓库、装依赖、传数据都不需要 GPU，只有编译 CUDA 算子和跑推理才需要。流程为：无卡开机配环境 → 关机 → 有卡开机跑任务 → 用完立即关机。
无卡模式 CPU 仅 0.5 核、内存 2 GB，编译 CUDA 算子会被 OOM 杀掉 （进程报 Killed），且编译器看不到 GPU、无法生成匹配显卡架构（如 3090 的 sm_86）的代码。因此 mamba 编译必须留到有卡模式。

1.2 环境配置命令

AutoDL 学术加速解决 GitHub 访问问题：

bash 复制代码

source /etc/network_turbo        # 开启学术加速
git clone https://github.com/TonyHongtaoWu/RainMamba.git

创建 conda 环境并安装 PyTorch（版本按项目发布年代配套，RainMamba 用 PyTorch 2.1.1 + CUDA 12.1 + Python 3.9）：

bash 复制代码

conda create -n rainmamba python=3.9 -y
source activate rainmamba
pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 \
    --index-url https://download.pytorch.org/whl/cu121

安装 MMEditing 依赖链（mim 会自动匹配预编译 wheel，避免源码编译 mmcv）：

bash 复制代码

pip install openmim timm
mim install mmcv-full==1.7.2
cd RainMamba/code/VRDS && pip install -e .    # 四个数据集目录各自独立，跑哪个装哪个

1.3 编译 mamba / causal-conv1d（有卡模式下）

bash 复制代码

pip install ninja==1.11.1.1
cd /root/RainMamba/causal-conv1d
MAX_JOBS=2 python setup.py install
cd ../mamba
MAX_JOBS=2 python setup.py install

MAX_JOBS=2 限制并行编译任务数，防止内存尖峰。mamba 核心文件 selective_scan_fwd.cu 单文件编译可达十几分钟，终端长时间"白字不动"属正常，可用 top 或 watch ls 观察编译产物确认存活。若需显式指定显卡架构（如 3090）：

bash 复制代码

TORCH_CUDA_ARCH_LIST="8.6" MAX_JOBS=2 python setup.py install

1.4 依赖版本踩坑记录

老项目复现的核心原则：依赖按项目发布年代配套，不装最新。本次实际遇到的四处版本冲突：

现象	原因	修复
`Killed` 编译中断	无卡模式内存不足 / OOM	切有卡模式 + `MAX_JOBS=2`
NumPy `_ARRAY_API not found`	pip 拉到 numpy 2.x，torch 2.1.1 用 NumPy 1.x 编译	`pip install numpy==1.26.4`
opencv 5.0 要求 numpy≥2	与上一步冲突	`pip install opencv-python==4.8.1.78`
einops 0.9.0.dev0 报 `TypeAlias` 不存在	开发版要求 Python≥3.10	`pip install einops==0.7.0`
transformers 4.57 报 `register_pytree_node` 缺失	版本过新，与 torch 2.1.1 不匹配	`pip install transformers==4.35.2`
`setup.py install` 装 tokenizers 失败	tar.gz 无 setup.py	`pip install einops transformers tokenizers` 补装

1.5 环境验证

bash 复制代码

python -c "import causal_conv1d, mamba_ssm; print('mamba ok')"

出现 mamba ok 即环境全通。

2. 预训练权重 VRDS.pth 下载

权重来源（README 提供的 Google Drive 链接，注意与"可视化结果"链接区分）：

VRDS.pth：drive.google.com/file/d/1iIkImRsHLYOdeSNvCRlm3SZvWilknF5c

服务器端直接下载尝试（需先 pip install gdown）：

bash 复制代码

source /etc/network_turbo
gdown 1iIkImRsHLYOdeSNvCRlm3SZvWilknF5c -O /root/RainMamba/code/VRDS/VRDS.pth

实测 gdown 报 "Cannot retrieve the public link"------Google Drive 对数据中心 IP 限流，属常见情况。稳定路线：本地浏览器（需代理）下载 → JupyterLab 上传至实例。450 MB 上传至 AutoDL 实测很快。

完整性校验：对比本地与服务器端文件字节数（ls -l 查看精确字节数），或两边各算一次 MD5（服务器 md5sum，Windows Get-FileHash -Algorithm MD5）比对。

另外 config 中 feat_pretrained 指向的 ConvNeXt-tiny 骨干权重（openmmlab 官方源，约 110 MB）无需手动下载，首次运行测试时自动拉取，国内直连可用。

3. 测试数据集下载与目录改造

3.1 下载

数据集源自 ViMP-Net（ACM MM'23）项目 OneDrive。VRDS 共 102 段视频（720p×100 帧），72 段训练 / 30 段测试；只复现推理只需要测试集（30 段、3000 帧）。OneDrive 目录下三个压缩包：

文件	大小	内容
testdrop.zip	2.37 GB	测试集序列（编号 000、001... 文件夹）
rainy.7z	4.46 GB	雨帧（对应 lq）
clean.7z	2.76 GB	干净帧（对应 gt）

判定某一包是 lq 还是 gt 的直接办法：打开一张 png，画面有雨丝/雨滴即 lq，画面干净即 gt。

3.2 目录结构改造

MMEditing 的 SRFolderMultipleGTDataset 要求 lq_folder 与 gt_folder 下各为一组同编号视频文件夹 ，内部帧文件名为 00000000.png 递增八位数字格式：

复制代码

data/VRDS/test/
├── lq/
│   ├── 000/00000000.png ... 00000099.png
│   ├── 001/...
│   └── ...（30 个）
└── gt/
    ├── 000/...
    └── ...

对应 config（configs/derainers/RainMamba/VRDS.py）中的相对路径：

python 复制代码

test=dict(
    type='SRFolderMultipleGTDataset',
    lq_folder="../data/VRDS/test/lq",
    gt_folder="../data/VRDS/test/gt",
    ...)

整理命令示例（假设解压出来是 rainy / clean 两个目录）：

bash 复制代码

mkdir -p /root/RainMamba/data/VRDS/test
cd /root/RainMamba/data/VRDS/test
mv rainy lq
mv clean gt

压缩包上传后先验证完整性再解压（避免解压中途才发现包损坏）：

bash 复制代码

unzip -t testdrop.zip      # 输出全 ok 即完好
unzip testdrop.zip

结构自检：

bash 复制代码

ls /root/RainMamba/data/VRDS/test/
ls /root/RainMamba/data/VRDS/test/lq/ | head -5

4. 推理：命令与代码

4.1 一条命令跑测试

bash 复制代码

cd /root/RainMamba/code/VRDS          # 必须在本目录下执行，config 中数据路径为相对路径
source activate rainmamba
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_test.sh \
    configs/derainers/RainMamba/VRDS.py "VRDS.pth" 1 \
    --save-path '../results/VRDS/' 2>&1 | tee test_log.txt

参数解读：

configs/derainers/RainMamba/VRDS.py：模型结构 + 数据路径配置文件；
"VRDS.pth"：加载的 checkpoint 路径（即权重文件位置，与执行目录相对）；
1：使用 GPU 数；
--save-path：复原结果输出目录；
tee test_log.txt：终端输出同步存档，便于事后 grep psnr test_log.txt。

4.2 这条命令背后发生了什么

dist_test.sh 内部实际调用 MMEditing 的测试入口 tools/test.py，流程为：

解析 config：按配置文件实例化模型（DrainNet 包裹 RainMamba 生成器）；
加载权重 ：从 .pth 恢复全部网络参数（load_checkpoint）；
构建测试数据集 ：扫描 lq/gt 目录，为每个视频序列生成索引（GenerateSegmentIndices）；测试时 num_input_frames 决定滑窗帧数（推理为时序模型，以滑动窗口方式逐段处理视频）；
前向推理 ：model(lq) 输出复原图像，无梯度计算（torch.no_grad）；
评估：逐帧与 gt 计算 PSNR / SSIM（crop_border=0），汇总打印。

4.3 结果查看与复现判据

数值指标 ：终端结尾输出 Evaluate conclusion: psnr: xx, ssim: xx；同时存在 test_log.txt；
复原结果 ：--save-path 指定目录（/root/RainMamba/results/VRDS/），每个测试视频一组输出帧，JupyterLab 双击可预览；
复现成功判据 ：对论文表 3，VRDS 测试集 PSNR 32.04 dB / SSIM 0.9366（ViMP-Net 为 31.02/0.9283）。差距 ±0.1 dB 内为完全复现；0.1--0.5 dB 一般为测试细节差异（帧对齐、裁边方式）；偏差过大需回查数据目录结构与权重完整性。

判断测试是否结束：终端回显 psnr 汇总并回到提示符即结束；辅助手段------另开终端 nvidia-smi 看 GPU 利用率归零，或 ls 输出目录看文件数停止增长。长任务建议 screen 挂后台防 SSH 断连。

4.4 demo：对任意雨天视频去雨

除跑满测试集外，仓库还提供单视频 demo（也吃 mp4）：

bash 复制代码

cd code/VRDS
python demo/restoration_video_demo.py \
    configs/derainers/RainMamba/VRDS.py "VRDS.pth" \
    input/rainy.mp4 output/derained.mp4

5. 成本复盘

按 6×3090 整机 ¥9.36/时计：

环节	耗时	费用
无卡模式装环境（torch/mmcv/mmedit）	1--2 h	≈ ¥0.2
有卡模式编译 mamba + causal-conv1d	0.5 h	≈ ¥5
权重/数据集上传（应切无卡）	0.5--1 h	≈ ¥0.1（无卡）
推理跑完 30 段测试视频	1--2 h	≈ ¥10--20
合计		≈ ¥15--25

优化空间：全程单卡 3090（¥1.32/时）+ 无卡模式传数据，总成本可压至 ¥5 以内。大文件传输务必切无卡模式，多卡机传一小时数据等于白烧 9 元。

6. 小结

复现推理的最短路径：代码（git clone）→ 权重（Google Drive/手机中转）→ 数据（OneDrive 下载、目录改造为 lq/gt 结构）→ dist_test.sh 一条命令出指标；
老项目复现的最大成本不在算法，在依赖版本考古：torch/mmcv/numpy/opencv/transformers 全部按项目年代配套；
权重是知识的实体化：作者 4 卡训练的成果以 450 MB 文件的形式被完整继承------这也是"模块缝合"式研究能成立的基础。

下一篇将解读 RainMamba 的网络结构（Hilbert 扫描机制与差分引导动态对比局部性学习），欢迎催更。