STCN

重新思考具有改进内存覆盖的时空网络,用于高效视频对象分割

Ho Kei Cheng, Yu-Wing Tai, Chi-Keung Tang

NeurIPS 2021

查看我们的新工作 Cutie!

bmx pigs

新闻: 在 YouTubeVOS 2021 挑战赛中,STCN 在新类别(未知类别)中获得第一名的准确率,在整体准确率中获得第二名。我们的解决方案也快速且轻量。

我们提出了时空对应网络(STCN)作为一个新的、有效的、高效的框架,用于在视频对象分割背景下建模时空对应关系。 STCN 在多个基准测试中达到了最先进的结果,同时运行速度快达20+ FPS,没有使用任何花哨的技巧。使用混合精度时,其速度甚至更高。尽管效果显著,但网络本身非常简单,还有很大的改进空间。有关技术细节请参阅论文。

更新 (2021年7月15日)

CBAM 模块:我们尝试了不使用 CBAM 模块,我认为我们确实不太需要它。对于 s03 模型,DAVIS 数据集上降低了 1.2,YouTubeVOS 数据集上提高了 0.1。对于 s012 模型,DAVIS 数据集上提高了 0.1,YouTubeVOS 数据集上提高了 0.1。欢迎您删除这个模块(参见 no_cbam 分支)。总的来说,规模更大的 YouTubeVOS 似乎是一个更好的一致性评估基准。

更新 (2021年8月22日)

可重复性:我们已更新了下面的包依赖要求。使用该环境,我们在两台不同的机器上进行了多次运行,在 DAVIS 数据集上获得了 [85.1, 85.5] 范围内的 J&F 分数。

更新 (2022年4月27日)

多尺度测试代码(如论文中所述)已添加在此。

这里有什么?

简单介绍
定量结果和预计算输出
1. DAVIS 2016
2. DAVIS 2017 验证集/测试开发集
3. YouTubeVOS 2018/2019
在您自己的数据上尝试我们的模型(提供交互式 GUI)
复现步骤
如果您想深入了解
引用

简单介绍

有两个主要贡献:STCN 框架(上图),和 L2 相似度。我们在图像之间而不是在(图像,掩码)对之间建立亲和力 - 这导致显著的速度提升、内存节省(因为我们只计算一个亲和力矩阵,而不是多个),和鲁棒性。我们进一步使用 L2 相似度来替代点积,这大大提高了内存库的利用率。

优点

简单,运行速度快(使用混合精度可达30+ FPS;不使用时可达20+)
高性能
仍有很大的改进空间(例如局部性、内存空间压缩)
易于训练:只需要两个11GB的GPU,不需要V100

要求

我们在开发这个项目时使用了以下软件包/版本。

PyTorch 1.8.1
torchvision 0.9.1
OpenCV 4.2.0
Pillow-SIMD 7.0.0.post3
progressbar2
thinspline 用于训练 (pip install git+https://github.com/cheind/py-thin-plate-spline)
gitpython 用于训练
gdown 用于下载预训练模型
[我的环境中的其他包](https://github.com/hkchengrex/STCN/blob/main/docs/packages.txt,仅供参考。

请参考官方 [PyTorch 指南](https://github.com/hkchengrex/STCN/blob/main/https://pytorch.org/ 安装 PyTorch/torchvision,以及 pillow-simd 指南安装 Pillow-SIMD。其余的可以通过以下命令安装:

pip install progressbar2 opencv-python gitpython gdown git+https://github.com/cheind/py-thin-plate-spline

结果

符号说明

FPS是摊销计算的,计算方法为总处理时间除以总帧数,与对象数量无关,即多对象FPS,在RTX 2080 Ti上测量,不包括IO时间。
我们还提供了使用自动混合精度(AMP)时的推理速度 -- 性能几乎相同。论文中的速度是在不使用AMP的情况下测量的。
所有评估都在480p分辨率下进行。为保持一致性,test-dev的FPS是在相同的内存设置下(每三帧作为记忆)在验证集上测量的。

[预计算输出 - Google Drive]

[预计算输出 - OneDrive]

s012表示经过BL预训练的模型,而s03表示未经预训练的模型(在MiVOS中曾称为s02)。

数据 (s012)

数据集	分割	J&F	J	F	FPS	FPS (AMP)
DAVIS 2016	验证	91.7	90.4	93.0	26.9	40.8
DAVIS 2017	验证	85.3	82.0	88.6	20.2	34.1
DAVIS 2017	test-dev	79.9	76.3	83.5	14.6	22.7

数据集	分割	总体得分	J-已见	F-已见	J-未见	F-未见
YouTubeVOS 18	验证	84.3	83.2	87.9	79.0	87.2
YouTubeVOS 19	验证	84.2	82.6	87.0	79.4	87.7

数据集	AUC-J&F	J&F @ 60s
DAVIS 交互式	88.4	88.8

对于DAVIS交互式,我们将MiVOS的传播模块从STM改为STCN。详情请参见此链接。

在自己的数据上尝试(提供交互式GUI)

如果你(以某种方式)有第一帧的分割(或更一般地,每个物体首次出现时的分割),你可以使用eval_generic.py。查看该文件顶部的说明。

如果你只是想交互式地玩一玩,我强烈推荐我们对MiVOS的扩展 :yellow_heart: -- 它带有交互式GUI,效率和效果都很高。

复现结果

预训练模型

我们对YouTubeVOS和DAVIS使用相同的模型。你可以自己下载并将它们放在./saves/中,或使用download_model.py。

s012模型(更好): [Google Drive] [OneDrive]

s03模型: [Google Drive] [OneDrive]

s0预训练模型: [GitHub]

s01预训练模型: [GitHub]

推理

eval_davis_2016.py 用于 DAVIS 2016 验证集
eval_davis.py 用于 DAVIS 2017 验证集和测试开发集（由 --split 控制）
eval_youtube.py 用于 YouTubeVOS 2018/19 验证集（由 --yv_path 控制）

参数提示应该能让你大致了解如何使用它们。例如，如果你已经使用我们的脚本下载了数据集和预训练模型，你只需要指定输出路径：python eval_davis.py --output [某个路径] 即可进行 DAVIS 2017 验证集评估。对于 YouTubeVOS 评估，将 --yv_path 指向你选择的版本。

多尺度测试代码（如论文中所述）已添加到这里。

训练

数据准备

我建议你可以软链接（ln -s）现有数据或使用提供的 download_datasets.py 来按照我们的格式构建数据集。download_datasets.py 可能会下载超出你需要的内容 —— 只需注释掉你不需要的部分。该脚本不会下载 BL30K，因为它非常大（>600GB），我们不想让你的硬盘崩溃。详见下文。

├── STCN
├── BL30K
├── DAVIS
│   ├── 2016
│   │   ├── Annotations
│   │   └── ...
│   └── 2017
│       ├── test-dev
│       │   ├── Annotations
│       │   └── ...
│       └── trainval
│           ├── Annotations
│           └── ...
├── static
│   ├── BIG_small
│   └── ...
├── YouTube
│   ├── all_frames
│   │   └── valid_all_frames
│   ├── train
│   ├── train_480p
│   └── valid
└── YouTube2018
    ├── all_frames
    │   └── valid_all_frames
    └── valid

BL30K

BL30K 是在 MiVOS 中提出的一个合成数据集。

你可以使用自动脚本 download_bl30k.py 或从 MiVOS 手动下载。请注意，每个片段大约 115GB —— 总共 700GB。你需要约 1TB 的可用磁盘空间来运行脚本（包括解压缓冲区）。 Google 可能会屏蔽 Google Drive 链接。你可以 1) 在你自己的 Google Drive 中为该文件夹创建快捷方式，2) 使用 rclone 从你自己的 Google Drive 复制（不会计入你的存储限制）。

训练命令

CUDA_VISIBLE_DEVICES=[a,b] OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port [cccc] --nproc_per_node=2 train.py --id [defg] --stage [h]

我们使用两个 11GB 的 GPU 实现了分布式数据并行（DDP）训练。将 a, b 替换为 GPU ID，cccc 替换为未使用的端口号，defg 替换为唯一的实验标识符，h 替换为训练阶段（0/1/2/3）。

模型通过不同的阶段进行渐进式训练（0：静态图像；1：BL30K；2：300K 主训练；3：150K 主训练）。每个阶段结束后，我们通过加载最新训练的权重来开始下一个阶段。

（仅在阶段 0 训练的模型不能直接使用。请参阅 model/model.py: load_network 了解我们所做的必要映射。）

带有 _checkpoint 后缀的 .pth 文件用于恢复中断的训练（使用 --load_model），通常不需要。一般情况下，你只需要 --load_network 并加载最后的网络权重（名称中不包含 checkpoint）。

因此，要训练 s012 模型，我们按顺序启动三个训练步骤，如下所示：

在静态图像上预训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s0 --stage 0

在 BL30K 数据集上预训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s01 --load_network [path_to_trained_s0.pth] --stage 1

主训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s012 --load_network [path_to_trained_s01.pth] --stage 2

要训练 s03 模型，我们按顺序启动两个训练步骤，如下所示：静态图像预训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s0 --stage 0

主要训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s03 --load_network [path_to_trained_s0.pth] --stage 3

深入了解

添加数据集或进行数据增强：dataset/static_dataset.py，dataset/vos_dataset.py
处理相似度函数或内存读取过程：model/network.py: MemoryReader，inference_memory_bank.py
处理网络结构：model/network.py，model/modules.py，model/eval_network.py
处理传播过程：model/model.py，eval_*.py，inference_*.py

引用

如果您发现本仓库有用，请引用我们的论文（如果使用top-k，请引用MiVOS）！

@inproceedings{cheng2021stcn,
  title={重新思考时空网络：改进内存覆盖以实现高效视频对象分割},
  author={Cheng, Ho Kei and Tai, Yu-Wing and Tang, Chi-Keung},
  booktitle={NeurIPS},
  year={2021}
}

@inproceedings{cheng2021mivos,
  title={模块化交互式视频对象分割：交互到蒙版、传播和差异感知融合},
  author={Cheng, Ho Kei and Tai, Yu-Wing and Tang, Chi-Keung},
  booktitle={CVPR},
  year={2021}
}

如果您想引用数据集： @inproceedings{shapenet2015, 标题 = {{ShapeNet：一个信息丰富的三维模型仓库}}, 作者 = {Chang, Angel Xuan 和 Funkhouser, Thomas 和 Guibas, Leonidas 和 Hanrahan, Pat 和 Huang, Qixing 和 Li, Zimo 和 Savarese, Silvio 和 Savva, Manolis 和 Song, Shuran 和 Su, Hao 和 Xiao, Jianxiong 和 Yi, Li 和 Yu, Fisher}, 会议论文集 = {arXiv:1512.03012}, 年份 = {2015} }