VPD 视觉感知框架实战:在 MMSegmentation 中用 Stable Diffusion 主干完成单目深度估计
【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation
导读
VPD(Visual Perception with a pre-trained Diffusion model)是一种将大规模图文预训练的文本到图像扩散模型(Stable Diffusion)直接用作视觉感知任务骨干网络的框架。本文以 configs/vpd/README.md 为骨架,结合 MMSegmentation 仓库中 VPD 的完整配置与源码实现,系统讲解该方法的原理、依赖安装、配置文件逐项解析、训练与推理流程,以及 NYU 数据集上的官方结果与评估指标。读完本文,你将掌握如何在 MMSegmentation 中复现、改造并部署基于 Stable Diffusion 主干的深度估计模型,并理解其背后的跨注意力引导与文本适配机制。
VPD 方法简介
论文出处:Unleashing Text-to-Image Diffusion Models for Visual Perception(ICCV 2023,官方仓库为 wl-zhao/VPD)
扩散模型(Diffusion Models)已成为生成式模型的新趋势,而基于大规模图文对预训练的文本到图像扩散模型可以通过提示词(prompt)进行高度可控的条件合成。与只关注底层属性和细节的无条件生成模型不同,文本到图像扩散模型得益于视觉-语言联合预训练,蕴含更丰富的高层语义知识。
VPD 框架的核心思想是:不把预训练去噪自编码器放在扩散采样流程中,而是直接把它当作骨干网络(backbone)使用,研究如何最大化利用其中学到的知识。具体做法有三点:
- 文本提示去噪解码器:用合适的文本输入(class embedding)去“提示”扩散模型的 UNet 解码器;
- 文本适配器(Text Adapter):通过可学习的适配器精炼文本特征,使其与预训练阶段对齐,并让视觉内容与文本提示充分交互;
- 跨注意力图显式引导:利用视觉特征与文本特征之间的跨注意力图(cross-attention maps)提供显式引导。
相比其他预训练方法,VPD 表明视觉-语言预训练的扩散模型可以借助该框架更快地适配下游视觉感知任务。论文在语义分割、指代图像分割(referring image segmentation)和深度估计三个任务上验证了有效性。当前 MMSegmentation 仓库中实现并开源的是其中的单目深度估计部分(基于 NYU 数据集)。
MMSegmentation 中的 VPD 实现架构
MMSegmentation 对 VPD 的实现分布在三个核心模块中,均通过注册器(@MODELS.register_module())挂载:
| 模块 | 文件 | 作用 |
|---|---|---|
VPD骨干网络 | mmseg/models/backbones/vpd.py | 加载 Stable Diffusion 预训练模型,提取多尺度视觉特征 |
VPDDepthHead解码头 | mmseg/models/decode_heads/vpd_depth_head.py | 将多尺度特征融合并回归出深度图 |
DepthEstimator分割器 | mmseg/models/segmentors/depth_estimator.py | 组织训练 / 推理 / 后处理全流程 |
VPD 骨干网络(mmseg/models/backbones/vpd.py)
VPD 骨干在 vpd.py 第 302 行 注册,其__init__主要完成四件事:
- 实例化并加载 Stable Diffusion 预训练模型:通过
instantiate_from_config(diffusion_cfg)构建完整的 Latent Diffusion 模型,再以非严格模式加载预训练权重(见 vpd.py 第 344-347 行); - 拆分出 VAE 编码器与 UNet:
self.encoder_vq = sd_model.first_stage_model(将 RGB 图像编码为 latent),self.unet = UNetWrapper(sd_model.model, **unet_cfg)(去噪 UNet 包装器,见 vpd.py 第 349-350 行); - 加载类别嵌入与文本适配器:从
class_embed_path加载 class embeddings,构建TextAdapter(一个Linear → GELU → Linear的残差适配模块,见 vpd.py 第 282-298 行),并以可学习参数gamma控制适配强度(默认1e-4); - 可选类别选择:当
class_embed_select=True时,会在 class embeddings 末尾拼接其均值向量(对应“无类别”情况),前向时根据每张图的category_id选择对应的嵌入(见 vpd.py 第 357-362 行)。
forward流程(见 vpd.py 第 365-395 行)可概括为:
输入图像 x ├─ 选择 class embeddings → TextAdapter(gamma) → 文本条件 c_crossattn ├─ 按 pad_shape 对输入做 padding(适配预训练扩散模型的输入尺寸) ├─ torch.no_grad() 下 encoder_vq.encode(x).mode() 得到 4 通道 latents ├─ unet(latents, t=1, context=c_crossattn) # 固定时间步,一次前向 └─ 返回多尺度特征列表(含跨注意力图拼接,经 UNetWrapper 逆序输出)值得注意的是,这里并不执行真正的扩散采样:时间步固定为t=1,encoder_vq前向全程torch.no_grad()(冻结 VAE),只有 UNet 及文本适配器参与梯度更新。UNetWrapper(vpd.py 第 193 行)在use_attn=True时通过register_attention_control(vpd.py 第 27 行)改写所有CrossAttention的前向逻辑,将跨注意力图经AttentionStore(vpd.py 第 107 行)收集、平均后,按分辨率(size16 / size32 / size64)拼接到 UNet 各层级输出上,从而把“视觉特征与文本提示的交互”显式注入到特征图中。
VPD 深度解码头(mmseg/models/decode_heads/vpd_depth_head.py)
VPDDepthHead(vpd_depth_head.py 第 100 行)接收骨干输出的四个层级特征,前向过程(见 vpd_depth_head.py 第 201-216 行):
- 将
x[3]上采样 2 倍后与x[2]拼接; - 对
x[0]、x[1]分别做 stride=2 的卷积下采样(conv1/conv2),使三者空间分辨率对齐; - 拼接后经
conv_aggregation(1x1 卷积 + GroupNorm + ReLU)聚合为feature_dim(默认 1536)维特征; - 按
fmap_border裁剪边缘后送入VPDDepthDecoder(3 层转置卷积 + 双线性上采样,见 vpd_depth_head.py 第 17 行); - 经
depth_pred_layer输出单通道,最后sigmoid * max_depth映射到实际深度范围(默认max_depth=10米)。
损失函数默认使用SiLogLoss(尺度不变对数损失),训练时会把预测深度 resize 回 GT 尺寸再计算损失(见 vpd_depth_head.py 第 218-253 行)。
深度估计器(mmseg/models/segmentors/depth_estimator.py)
DepthEstimator(depth_estimator.py 第 21 行)继承自EncoderDecoder,负责组织完整的训练 / 推理生命周期:
loss():extract_feat → decode_head.loss;predict():inference → postprocess_result,结果写入SegDataSample.pred_depth_map(见 depth_estimator.py 第 333-391 行);- 支持
whole/slide/slide_flip三种推理模式,其中slide_flip会以滑窗方式推理并对原图与水平翻转结果取平均(见 depth_estimator.py 第 241-298 行),VPD 配置默认使用该模式以提升精度; extract_feat在class_embed_select=True时会把batch_img_metas中的category_id组装成类别标签传给骨干(见 depth_estimator.py 第 98-113 行)。
仓库对应单元测试位于 tests/test_models/test_backbones/test_vpd.py、tests/test_models/test_heads/test_vpd_depth_head.py 与 tests/test_models/test_segmentors/test_depth_estimator.py,可用于验证各模块的输入输出尺寸与基本前向。
环境依赖安装
VPD 依赖ldm(Latent Diffusion Models)库与 albumentations 增强库,运行前需要按 README 安装:
pip install -r requirements/albu.txt pip install -r requirements/optional.txt其中:
- requirements/albu.txt 提供 Albumentations 数据增强相关依赖,对应 NYU 训练管线中使用的
Albu变换; - requirements/optional.txt 提供 LDM 依赖;若未安装,
VPD骨干与UNetWrapper在初始化时会抛出断言错误,提示先安装该依赖(见 vpd.py 第 214-215 行 与 vpd.py 第 333-334 行)。
VPD 的完整模型、权重与日志元信息(许可证、训练数据、指标、权重与日志下载地址等)记录在 configs/vpd/metafile.yaml 中,可配合tools/model_converters与mim工具链使用。
配置文件逐项解析
VPD 在 MMSegmentation 中共有两个训练配置,均以“SD-v1.5 骨干 + NYU 数据集 + 25k 迭代”为基线,区别仅在于裁剪尺寸(480×480 与 512×512):
- configs/vpd/vpd_sd_4xb8-25k_nyu-480x480.py
- configs/vpd/vpd_sd_4xb8-25k_nyu-512x512.py
配置继承关系
两个配置都继承自四份基础配置:
_base_ = [ '../_base_/models/vpd_sd.py', '../_base_/datasets/nyu.py', '../_base_/default_runtime.py', '../_base_/schedules/schedule_25k.py' ]- configs/base/models/vpd_sd.py:模型结构 + Stable Diffusion 完整结构描述;
- configs/base/datasets/nyu.py:NYU 数据加载与增强管线;
- configs/base/default_runtime.py:运行时钩子、日志、环境等;
- configs/base/schedules/schedule_25k.py:优化器与学习率调度。
模型主干配置(configs/base/models/vpd_sd.py)
该文件定义了完整的 Stable Diffusion 结构(适配自官方v1-inference.yaml),关键参数包括:
| 参数 | 值 | 说明 |
|---|---|---|
base_learning_rate | 0.0001 | 预训练基础学习率(本仓库训练时会被 schedule 覆盖) |
timesteps | 1000 | 扩散时间步总数 |
linear_start/linear_end | 0.00085 / 0.012 | latent 扩散噪声调度区间 |
scale_factor | 0.18215 | VAE latent 缩放系数 |
image_size/channels | 64 / 4 | latent 空间尺寸与通道数 |
cond_stage_trainable | False | 文本编码器不参与训练 |
conditioning_key | crossattn | 使用跨注意力条件注入 |
UNetmodel_channels | 320 | UNet 基础通道数 |
UNetchannel_mult | [1, 2, 4, 4] | 四阶段通道倍增 |
UNetcontext_dim | 768 | 文本条件维度(对应 CLIP 文本编码) |
VAEembed_dim | 4 | 自编码器 latent 维度 |
模型部分组装为DepthEstimator,并设置find_unused_parameters = True(Stable Diffusion 中有部分参数不参与梯度更新,需开启该标志避免 DDP 报错):
model = dict( type='DepthEstimator', data_preprocessor=data_preprocessor, backbone=dict( type='VPD', diffusion_cfg=stable_diffusion_cfg, ), )data_preprocessor使用SegDataPreProcessor,mean=[127.5, 127.5, 127.5]、std=[127.5, 127.5, 127.5],即把输入归一化到[-1, 1],与 Stable Diffusion 预训练时的数据分布保持一致。
下游任务配置(以 480×480 为例)
crop_size = (480, 480) model = dict( type='DepthEstimator', data_preprocessor=dict(size=crop_size), backbone=dict( class_embed_path='https://download.openmmlab.com/mmsegmentation/' 'v0.5/vpd/nyu_class_embeddings.pth', class_embed_select=True, pad_shape=512, unet_cfg=dict(use_attn=False), ), decode_head=dict( type='VPDDepthHead', in_channels=[320, 640, 1280, 1280], max_depth=10, fmap_border=(1, 1), ), test_cfg=dict(mode='slide_flip', crop_size=crop_size, stride=(160, 160)))各关键项说明:
class_embed_path:NYU 数据集的类别嵌入权重下载地址(运行时会自动下载),用于把“类别”作为文本提示注入 UNet。class_embed_select=True表示启用类别选择机制,配合数据集 meta 中的category_id使用;pad_shape=512:输入图像先 pad 到 512,以满足预训练扩散模型对输入尺寸的要求(实现见 vpd.py 第 384-387 行);unet_cfg=dict(use_attn=False):深度估计任务关闭跨注意力图收集与拼接(该机制主要服务于分割类任务),减少显存与计算开销;decode_head:VPDDepthHead,四层级输入通道[320, 640, 1280, 1280],max_depth=10(预测深度上限 10 米),480 分辨率下设置fmap_border=(1, 1)裁剪 1 像素边缘以对齐尺寸;test_cfg:推理采用slide_flip滑窗 + 翻转平均;480 配置 stride 为(160, 160),512 配置 stride 为(128, 128)。
检查点与优化器定制
两个配置都通过default_hooks按 RMSE 保存最优权重:
default_hooks = dict( checkpoint=dict(save_best='rmse', rule='less', max_keep_ckpts=1))优化器部分使用自定义构造器ForceDefaultOptimWrapperConstructor,对不同子网络设置差异化学习率:
optim_wrapper = dict( constructor='ForceDefaultOptimWrapperConstructor', paramwise_cfg=dict( bias_decay_mult=0, force_default_settings=True, custom_keys={ 'backbone.encoder_vq': dict(lr_mult=0), 'backbone.unet': dict(lr_mult=0.01), }))backbone.encoder_vq(VAE 编码器)lr_mult=0:完全冻结,只做前向编码,对应源码中torch.no_grad()的实现;backbone.unetlr_mult=0.01:UNet 使用主学习率的 1%(微调);bias_decay_mult=0:偏置项不做权重衰减。
基础优化器与调度定义在 configs/base/schedules/schedule_25k.py:
optimizer = dict(type='AdamW', lr=0.001, weight_decay=0.1) param_scheduler = [ dict(type='LinearLR', start_factor=3e-2, begin=0, end=12000, by_epoch=False), dict(type='PolyLRRatio', eta_min_ratio=3e-2, power=0.9, begin=12000, end=24000, by_epoch=False), dict(type='ConstantLR', by_epoch=False, factor=1, begin=24000, end=25000) ] train_cfg = dict(type='IterBasedTrainLoop', max_iters=25000, val_interval=1000)即:前 12k 迭代线性热身(起始为 3e-2),随后多项式衰减至 24000 步,最后 1k 步保持恒定,共训练 25000 迭代,每 1000 迭代验证一次。
NYU 数据管线(configs/base/datasets/nyu.py)
NYU 数据集默认根目录为data/nyu,训练图像位于images/train,深度标注位于annotations/train。训练管线要点:
LoadDepthAnnotation:加载深度图并乘以depth_rescale_factor=1e-3(把毫米量级的原始深度换算为米);RandomDepthMix(prob=0.25):VPD 配套的深度随机混合增强;RandomFlip(prob=0.5)+RandomCrop(crop_size=(480, 480));Albu增强:RandomBrightnessContrast、RandomGamma、HueSaturationValue(对应requirements/albu.txt);PackSegInputs的meta_keys中包含category_id,供class_embed_select使用。
验证/测试管线会将图像按scale=(2000, 480)等比缩放,评估器为DepthMetric:
val_evaluator = dict( type='DepthMetric', min_depth_eval=0.001, max_depth_eval=10.0, crop_type='nyu_crop')即只在 0.001~10 米有效深度范围内、按 NYU 标准裁剪区域评估。
训练与推理
训练
准备好 NYU 数据集(放置于data/nyu,结构符合 configs/base/datasets/nyu.py 中的data_prefix)后,使用仓库统一的训练入口 tools/train.py 启动:
# 单卡训练 480x480 配置 python tools/train.py configs/vpd/vpd_sd_4xb8-25k_nyu-480x480.py # 多卡分布式训练(8 卡) bash tools/dist_train.sh configs/vpd/vpd_sd_4xb8-25k_nyu-480x480.py 8训练过程中,UNet 以 1% 学习率微调、VAE 完全冻结、文本适配器与解码头正常更新;训练前会自动下载 Stable Diffusion 预训练权重与nyu_class_embeddings.pth。训练迭代数较大且显存占用较高,建议参考 metafile 中“8× A100、batch size 32”的资源前提按需调整。
推理 / 测试
使用 tools/test.py 配合权重文件进行测试,结果会按slide_flip模式输出深度图并计算 RMSE 等指标:
python tools/test.py configs/vpd/vpd_sd_4xb8-25k_nyu-480x480.py ${CHECKPOINT_FILE}模型权重与训练日志的下载地址均记录在 configs/vpd/metafile.yaml 中(Weights与Training log字段),可参照其中文件名组织本地权重路径。
官方结果与评估指标
NYU 深度估计结果
以下是 README 与 metafile 中记录的两个 VPD 模型在 NYU 数据集上的官方指标(训练资源:8× A100,batch size 32):
| 方法 | 骨干网络 | Crop Size | Lr schd | 设备 | RMSE | d1 | d2 | d3 | REL | log_10 | 配置 |
|---|---|---|---|---|---|---|---|---|---|---|---|
| VPD | Stable-Diffusion-v1-5 | 480x480 | 25000 | A100 | 0.253 | 0.964 | 0.995 | 0.999 | 0.069 | 0.030 | config |
| VPD | Stable-Diffusion-v1-5 | 512x512 | 25000 | A100 | 0.258 | 0.963 | 0.995 | 0.999 | 0.072 | 0.031 | config |
指标含义(数值越小越好的为 RMSE、REL、log_10;越接近 1 越好的是 d1/d2/d3):
- RMSE:深度预测的均方根误差;
- REL:平均相对误差(
|pred - gt| / gt的均值); - log_10:以 10 为底的对数空间误差;
- d1 / d2 / d3:阈值准确率,即预测与真值比值落在
[1/1.25, 1.25]、[1/1.25², 1.25²]、[1/1.25³, 1.25³]内的像素比例。
论文指出,VPD 在 NYUv2 深度估计上达到 0.254 RMSE,并在 RefCOCO-val 指代分割上达到 73.3% oIoU,均刷新了当时的榜单纪录。MMSegmentation 中的 480×480 配置复现结果 RMSE 为 0.253,与论文结论一致。
引用
若在研究中使用了 VPD 方法,请按 README 给出的文献引用:
@article{zhao2023unleashing, title={Unleashing Text-to-Image Diffusion Models for Visual Perception}, author={Zhao, Wenliang and Rao, Yongming and Liu, Zuyan and Liu, Benlin and Zhou, Jie and Lu, Jiwen}, journal={ICCV}, year={2023} }小结
VPD 为“如何复用扩散模型的语义知识”提供了一个简洁而有效的范式:冻结 VAE、微调 UNet、以类别嵌入作为文本提示、用可学习适配器对齐图文特征。在 MMSegmentation 中,它被完整落地为VPD骨干 +VPDDepthHead解码头 +DepthEstimator分割器的标准三段式结构,开箱即用地支持训练、测试与指标评估。读者可以从 configs/vpd/vpd_sd_4xb8-25k_nyu-480x480.py 出发,复现官方结果;也可参考 mmseg/models/backbones/vpd.py 与 mmseg/models/decode_heads/vpd_depth_head.py 的源码,将该范式迁移到语义分割、指代分割等更多视觉感知任务上。
【免费下载链接】mmsegmentationOpenMMLab Semantic Segmentation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmsegmentation
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考