Segmenter：用于语义分割的Transformer

Segmenter：用于语义分割的Transformer
作者：Robin Strudel*、Ricardo Garcia*、Ivan Laptev 和 Cordelia Schmid，ICCV 2021。

*贡献相等

🔥 Segmenter 现已在 MMSegmentation 中可用。

安装

在你的 .bashrc 文件中定义指向检查点和数据集目录的环境变量：

export DATASET=/path/to/dataset/dir

安装 PyTorch 1.9，然后在本仓库的根目录下运行 pip install .。

要下载 ADE20K 数据集，使用以下命令：

python -m segm.scripts.prepare_ade20k $DATASET

模型库

我们发布了以 Vision Transformer 为骨干网络的模型，其权重初始化自 改进版 ViT 模型。

ADE20K

带有 ViT 骨干网络的 Segmenter 模型：

名称	mIoU (SS/MS)	# 参数	分辨率	FPS	下载
Seg-T-Mask/16	38.1 / 38.8	7M	512x512	52.4	模型	配置文件	日志
Seg-S-Mask/16	45.3 / 46.9	27M	512x512	34.8	模型	配置文件	日志
Seg-B-Mask/16	48.5 / 50.0	106M	512x512	24.1	模型	配置文件	日志
Seg-B/8	49.5 / 50.5	89M	512x512	4.2	模型	配置文件	日志
Seg-L-Mask/16	51.8 / 53.6	334M	640x640	-	模型	配置文件	日志

带有 DeiT 骨干网络的 Segmenter 模型：

名称	mIoU (SS/MS)	# 参数	分辨率	FPS	下载
Seg-B†/16	47.1 / 48.1	87M	512x512	27.3	模型	配置文件	日志
Seg-B†-Mask/16	48.7 / 50.1	106M	512x512	24.1	模型	配置文件	日志

Pascal Context

名称	mIoU (SS/MS)	# 参数	分辨率	FPS	下载
Seg-L-Mask/16	58.1 / 59.0	334M	480x480	-	模型	配置文件	日志

Cityscapes

名称	mIoU (SS/MS)	# 参数	分辨率	FPS	下载
Seg-L-Mask/16	79.1 / 81.3	322M	768x768	-	模型	配置文件	日志

推理

将一个检查点及其配置文件下载到同一个文件夹中，例如 seg_tiny_mask。

你可以使用以下命令从自己的数据生成分割图：

python -m segm.inference --model-path seg_tiny_mask/checkpoint.pth -i images/ -o segmaps/ 

要在 ADE20K 上进行评估，运行以下命令：

# 单尺度评估：
python -m segm.eval.miou seg_tiny_mask/checkpoint.pth ade20k --singlescale
# 多尺度评估：
python -m segm.eval.miou seg_tiny_mask/checkpoint.pth ade20k --multiscale

训练

在单个 GPU 上对 ADE20K 数据集训练 Seg-T-Mask/16：

python -m segm.train --log-dir seg_tiny_mask --dataset ade20k \
  --backbone vit_tiny_patch16_384 --decoder mask_transformer

要训练 Seg-B-Mask/16，只需将骨干网络设置为 vit_base_patch16_384，并使用至少 4 张 V100 GPU（每轮大约 12 分钟）或最多 8 张 V100 GPU（每轮大约 7 分钟）运行上述命令。代码使用 SLURM 环境变量。

日志

要绘制实验日志，可以使用以下命令：

python -m segm.utils.logs logs.yml

其中 logs.yml 位于 utils/ 目录下，内容如下，指定了您的实验日志路径：

root: /path/to/checkpoints/
logs:
  seg-t: seg_tiny_mask/log.txt
  seg-b: seg_base_mask/log.txt

注意力图

要可视化 Seg-T-Mask/16 编码器第 0 层中位置为 (0, 21) 的 patch 的注意力图，可以使用以下命令：

python -m segm.scripts.show_attn_map seg_tiny_mask/checkpoint.pth \ 
images/im0.jpg output_dir/ --layer-id 0 --x-patch 0 --y-patch 21 --enc

提供了不同的选项来选择生成的注意力图：

• --enc 或 --dec：分别选择编码器或解码器的注意力图。
• --patch 或 --cls：--patch 会生成对应坐标为 (x_patch, y_patch) 的 patch 的注意力图。--cls 结合 --enc 可以生成编码器 CLS token 的注意力图，而结合 --dec 则会生成解码器每个类别嵌入的注意力图。
• --x-patch 和 --y-patch：用于指定绘制注意力图的 patch 坐标。当使用 --cls 时，此参数将被忽略。
• --layer-id：选择生成注意力图的层号。

例如，要生成解码器类别嵌入的注意力图，可以使用以下命令：

python -m segm.scripts.show_attn_map seg_tiny_mask/checkpoint.pth \
images/im0.jpg output_dir/ --layer-id 0 --dec --cls

Seg-L-Mask/16 编码器第 1、4、8、12 和 16 层中位置为 (0, 21) 的 patch 的注意力图：

Seg-L-Mask/16 解码器第 0 层中类别嵌入的注意力图：

视频分割

使用在 ADE20K 上训练的 Seg-B-Mask/16 模型，对 DAVIS 视频数据集进行零样本视频分割。

BibTex

@article{strudel2021,
  title={Segmenter: Transformer for Semantic Segmentation},
  author={Strudel, Robin and Garcia, Ricardo and Laptev, Ivan and Schmid, Cordelia},
  journal={arXiv preprint arXiv:2105.05633},
  year={2021}
}

致谢

Vision Transformer 的代码基于 timm 库，语义分割的训练和评估流程则使用了 mmsegmentation。

Segmenter 快速上手指南

Segmenter 是一个基于 Transformer 的语义分割模型，由 INRIA 团队提出。本指南将帮助你快速在本地环境中部署并运行该工具。

1. 环境准备

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

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.x
• 深度学习框架: PyTorch 1.9+
• 硬件: 支持 CUDA 的 NVIDIA GPU（训练大模型建议多卡环境）
• 依赖库: timm (用于 Vision Transformer 骨干网络)

注意：本项目代码基于 timm 和 MMSegmentation 的部分逻辑构建。

2. 安装步骤

2.1 配置环境变量

首先，定义数据集和检查点（模型权重）的存放路径。请将以下命令添加到你的 ~/.bashrc 文件中，或直接并在当前终端执行：

export DATASET=/path/to/dataset/dir

(请将 /path/to/dataset/dir 替换为你实际的磁盘路径)

2.2 安装 PyTorch

访问 PyTorch 官网 获取适合你环境的安装命令。例如，使用 pip 安装 PyTorch 1.9：

pip install torch==1.9.0 torchvision==0.10.0 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

(国内用户可使用清华源加速：pip install ... -i https://pypi.tuna.tsinghua.edu.cn/simple)

2.3 安装 Segmenter

克隆仓库或在仓库根目录下执行以下命令进行安装：

pip install .

2.4 准备数据集 (以 ADE20K 为例)

运行以下脚本自动下载并预处理 ADE20K 数据集：

python -m segm.scripts.prepare_ade20k $DATASET

3. 基本使用

3.1 下载预训练模型

从 Model Zoo 中选择一个模型（例如 Seg-T-Mask/16），下载其 checkpoint.pth 和 variant.yml 配置文件到同一文件夹（例如 seg_tiny_mask）。

• 模型权重: 下载链接
• 配置文件: 下载链接

3.2 图像推理 (Inference)

使用下载的模型对你的图片进行语义分割。假设你的输入图片在 images/ 目录，输出结果保存到 segmaps/：

python -m segm.inference --model-path seg_tiny_mask/checkpoint.pth -i images/ -o segmaps/ 

3.3 模型评估 (Evaluation)

如果你已准备好 ADE20K 验证集，可以运行以下命令评估模型的 mIoU 指标：

单尺度评估：

python -m segm.eval.miou seg_tiny_mask/checkpoint.pth ade20k --singlescale

多尺度评估：

python -m segm.eval.miou seg_tiny_mask/checkpoint.pth ade20k --multiscale

3.4 开始训练 (Training)

在单张 GPU 上训练 Seg-T-Mask/16 模型：

python -m segm.train --log-dir seg_tiny_mask --dataset ade20k \
  --backbone vit_tiny_patch16_384 --decoder mask_transformer

提示：若训练更大的模型（如 Seg-B-Mask/16），建议修改 --backbone 为 vit_base_patch16_384，并使用至少 4 张 V100 GPU 进行分布式训练。