论文地址:https://arxiv.org/pdf/2312.02145
源码地址:https://github.com/prs-eth/marigold
概述
传统的单目深度估计模型往往在处理含有不熟悉内容和布局的图像时表现不佳,因为它们的视觉知识受限于训练数据,尤其在面对新领域的零样本(zero-shot)泛化时更是如此。尽管当前的单目深度估计模型已经从CNN发展到大型视觉Transformer,但依然存在泛化能力不足的问题。为此,作者提出了一个新的方法,名为 Marigold,它利用了现代生成式扩散模型中丰富的视觉先验知识,特别是基于Stable Diffusion模型的改进。Marigold的核心贡献在于:
- 提出了一种简单且资源高效的微调协议,将预训练的Latent Diffusion Model(LDM)转化为图像条件的深度估计模型,同时保留其丰富的视觉先验知识。
- Marigold在多种数据集上实现了最先进的单目深度估计性能,展示了优异的零样本泛化能力,特别是在某些情况下取得了显著的性能提升。
通过使用合成的RGB-D数据集对Marigold进行微调,该模型能够在单个GPU上快速训练,且无需在新数据上进行广泛的再训练即可应对多样且真实的场景。
方法
将单目深度估计问题表述为一个条件去噪扩散生成任务。具体步骤如下:
-
前向过程(Forward Process):从条件分布开始 d 0 : = d d_0 := d d0:=d,逐步添加高斯噪声,在 t ∈ { 1 , . . . , T } t \in \{1, ..., T\} t∈{1,...,T}的每个噪声等级下获得噪声样本 d t d_t dt:
d t = α ˉ t d 0 + 1 − α ˉ t ϵ (1) d_t = \sqrt{\bar{\alpha}_t} d_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon\tag{1} dt=αˉtd0+1−αˉtϵ(1)
其中, ϵ ∼ N ( 0 , I ) \epsilon \sim \mathcal{N}(0, I) ϵ∼N(0,I), α ˉ t = ∏ s = 1 t ( 1 − β s ) \bar{\alpha}_t = \prod_{s=1}^{t} (1 - \beta_s) αˉt=∏s=1t(1−βs), { β 1 , … , β T } \{\beta_1, \ldots, \beta_T\} {β1,…,βT} 是扩散过程的方差调度。
-
反向过程(Reverse Process):利用学习到的去噪模型 ϵ θ ( ⋅ ) \epsilon_\theta(\cdot) ϵθ(⋅) 逐步移除 d t d_t dt 中的噪声,以获得 d t − 1 d_{t-1} dt−1。模型的参数 θ \theta θ 通过以下过程进行更新:
- 从训练集中获取一个数据对 ( x , d ) (x, d) (x,d);
- 在随机时间步 t t t,对 d d d 添加噪声 ϵ \epsilon ϵ;
- 计算噪声估计 ϵ ^ = ϵ θ ( d t , x , t ) \hat{\epsilon} = \epsilon_\theta(d_t, x, t) ϵ^=ϵθ(dt,x,t);
- 最小化去噪扩散目标函数:
L = E d 0 , ϵ ∼ N ( 0 , I ) , t ∼ U ( T ) [ ∥ ϵ − ϵ ^ ∥ 2 2 ] (2) L = \mathbb{E}_{d_0, \epsilon \sim \mathcal{N}(0, I), t \sim U(T)} \left[ \lVert \epsilon - \hat{\epsilon} \rVert^2_2 \right]\tag{2} L=Ed0,ϵ∼N(0,I),t∼U(T)[∥ϵ−ϵ^∥22](2)
-
推理过程(Inference Process):推理时,从正态分布变量 d T d_T dT 开始,通过迭代应用学习到的去噪模型 ϵ θ ( d t , x , t ) \epsilon_\theta(d_t, x, t) ϵθ(dt,x,t) 逐步重建 d = d 0 d = d_0 d=d0。以上过程将单目深度估计转化为一个条件去噪扩散生成任务,并通过调整去噪U-Net模型实现从噪声到深度图的生成。
训练阶段
文中基于Stable Diffusion V2的预训练框架,在潜在空间进行扩散训练。Marigold使用了冻结的VAE(变分自编码器)来对输入的图像和其对应的深度图进行编码。对于给定的深度图
d
d
d,沿着特征通道复制3份来模拟RGB图像输入。此外,将深度图
d
d
d 进行线性归一化到
[
−
1
,
1
]
[-1, 1]
[−1,1],以匹配VAE的设计输入值范围并确保了深度表示在不同数据统计条件下保持一致的仿射不变性:
d
~
=
(
d
−
d
2
%
d
98
%
−
d
2
%
−
0.5
)
×
2
(3)
\tilde{d} = \left(\frac{d - d_{2\%}}{d_{98\%} - d_{2\%}} - 0.5\right) \times 2\tag{3}
d~=(d98%−d2%d−d2%−0.5)×2(3)
其中,
d
2
%
d_{2\%}
d2% 和
d
98
%
d_{98\%}
d98% 分别表示深度图中第2百分位和第98百分位的深度值。该过程确保了场景的深度值在近平面和远平面之间得到合理的映射,使得模型能够聚焦于纯粹的仿射不变深度估计。
其对应的潜变量编码表示为 z ( d ) = E ( d ~ ) z^{(d)} = \mathcal{E}(\tilde{d}) z(d)=E(d~)。条件图像 x x x 也同样被转化到潜变量空间,表示为 z ( x ) = E ( x ) z^{(x)} = \mathcal{E}(x) z(x)=E(x)。
将图像和深度潜变量编码沿特征维度连接 z t = cat ( z t ( d ) , z ( x ) ) z_t = \text{cat}(z^{(d)}_t, z^{(x)}) zt=cat(zt(d),z(x))。为了适应扩展后的输入 z t z_t zt,Unet的输入通道数增加了一倍。为了防止输入层激活幅度的膨胀,并尽可能维持预训练结构不必拿,文中将输入层的权重张量直接复制扩展一倍,并将其值除以2。去噪器在潜变量空间中进行训练: ϵ θ ( z t ( d ) , z ( x ) , t ) \epsilon_\theta(z^{(d)}_t, z^{(x)}, t) ϵθ(zt(d),z(x),t)。
最后由解码器 D D D 从 z 0 ( d ) z^{(d)}_0 z0(d) 重建数据 d ^ \hat{d} d^: d ^ = D ( z 0 ( d ) ) \hat{d} = D(z^{(d)}_0) d^=D(z0(d)),最后将三通道的平均值作为预测的深度图。
训练数据
训练过程使用合成数据进行,因为合成深度图数据在每个像素上都具有有效的深度值,这使得数据更为完整且干净。此外,合成数据避免了真实深度数据中由于传感器物理特性而产生的噪声和丢失像素问题。通过这种方法,Marigold能够在真实场景中实现出色的零样本迁移性能。
多分辨率噪声
文中提出了退火多分辨率噪声的概念,通过在扩散过程中逐步减弱低分辨率噪声的影响,将噪声分布逐渐过渡到标准高斯噪声。这种方法能够提高模型的训练效率,并改善推理阶段的预测性能。退火策略通过逐步调整噪声层的权重,使得模型能够更好地适应不同尺度的噪声,有助于提高深度估计的精度和稳定性。
多分辨率噪声的生成 多分辨率噪声的生成过程如下:首先,从不同尺度采样多个高斯噪声图像,形成一个金字塔结构的噪声集合,然后对这些噪声图像进行上采样,使它们达到U-Net输入的分辨率。这些不同尺度的噪声图像通过加权平均和重新归一化组合成一个噪声输入。第
i
i
i 层金字塔的噪声权重为
s
i
s^i
si,其中
0
<
s
<
1
0 < s < 1
0<s<1 表示低分辨率噪声的影响强度。
退火策略 为了使多分辨率噪声逐渐逼近扩散模型中使用的标准高斯噪声,文中提出了一种退火策略。这种策略在扩散调度的每个时间步 t t t 上,对高于0级的噪声层赋予退火权重 ( s ⋅ t / T ) i (s \cdot t / T)^i (s⋅t/T)i,其中 T T T 是总扩散步数。这意味着在靠近无噪声终点的时间步,低分辨率层的噪声影响力会逐渐减少,从而逐步过渡到标准的高斯噪声。
性能提升 通过实验发现,使用这种退火的多分辨率噪声能够加速模型的收敛,并在推理时显著提高深度预测的质量。该策略特别适用于合成数据的短期微调过程,有助于模型更好地应对不同分辨率的噪声环境,提高最终的深度估计精度。
推理阶段
- 编码图像:首先,使用原始的Stable Diffusion VAE对输入图像 x x x 进行编码,得到图像的潜变量编码 z ( x ) = E ( x ) z^{(x)} = E(x) z(x)=E(x)。
- 初始化深度潜变量:将深度潜变量 z T ( d ) z^{(d)}_T zT(d) 初始化为标准高斯噪声(训练时采用多分辨率噪声)。
- 逐步去噪:在每个去噪步骤中,将当前深度潜变量 z t ( d ) z^{(d)}_t zt(d) 和条件图像潜变量 z ( x ) z^{(x)} z(x) 作为输入,使 U-Net 进行逆扩散过程。
- 非马尔可夫采样:作者采用了DDIM的方法进行非马尔可夫采样,并使用重采样步骤以加速推理。
- 解码潜变量:最终,通过解码器 D D D 对去噪后的潜变量 z 0 ( d ) z^{(d)}_0 z0(d) 进行解码,得到最终的深度图预测 d ^ = D ( z 0 ( d ) ) \hat{d} = D(z^{(d)}_0) d^=D(z0(d)),并对输出的三通道取平均值作为最终的深度图估计。
测试阶段集成
通过对同一输入进行多次推理来提升深度估计的稳定性和精度。由于Marigold模型的推理过程具有随机性,不同的初始噪声会导致预测结果有所差异,因此测试时集成可以有效地降低这种不确定性。
对于每个输入样本,模型可以执行多次推理。每次推理时,从不同的初始噪声
z
T
(
d
)
z^{(d)}_T
zT(d) 开始,得到一组仿射不变的深度预测
{
d
^
1
,
d
^
2
,
…
,
d
^
N
}
\{\hat{d}_1, \hat{d}_2, \ldots, \hat{d}_N\}
{d^1,d^2,…,d^N}。为了对这些预测进行聚合,作者提出了以下的集成方案:
-
对齐尺度和偏移:首先,对于每次推理得到的深度预测 d ^ i \hat{d}_i d^i,计算相对于某个标准尺度和范围的对应尺度 s i s_i si 和偏移 t i t_i ti,即 d ^ i ′ = d ^ i × s i + t i \hat{d}_i' = \hat{d}_i \times s_i + t_i d^i′=d^i×si+ti。
-
联合优化:通过最小化每对对齐后的预测 ( d ^ i ′ , d ^ j ′ ) (\hat{d}_i', \hat{d}_j') (d^i′,d^j′) 之间的距离来进行联合优化。具体目标函数为:
min s 1 , … , s N , t 1 , … , t N ( 1 b ∑ i = 1 N − 1 ∑ j = i + 1 N ∥ d ^ i ′ − d ^ j ′ ∥ 2 2 + λ R ) (4) \min_{s_1, \ldots, s_N, t_1, \ldots, t_N} \left( \sqrt{\frac{1}{b} \sum_{i=1}^{N-1} \sum_{j=i+1}^{N} \lVert \hat{d}_i' - \hat{d}_j' \rVert_2^2} + \lambda R \right)\tag{4} s1,…,sN,t1,…,tNmin b1i=1∑N−1j=i+1∑N∥d^i′−d^j′∥22+λR (4)其中, b = ( N 2 ) b = \binom{N}{2} b=(2N) 表示预测对的总数,正则化项 R = ∣ min ( m ) ∣ + ∣ 1 − max ( m ) ∣ R = |\min(m)| + |1 - \max(m)| R=∣min(m)∣+∣1−max(m)∣ 用于防止结果塌陷到一个平凡解,并强制输出单位尺度。
-
最终集成预测:在迭代优化对齐过程后,通过像素级中值计算得到集成后的深度图 m m m,公式为:
m ( x , y ) = median ( d ^ 1 ′ ( x , y ) , … , d ^ N ′ ( x , y ) ) (5) m(x, y) = \text{median}(\hat{d}_1'(x, y), \ldots, \hat{d}_N'(x, y))\tag{5} m(x,y)=median(d^1′(x,y),…,d^N′(x,y))(5)
该集成过程不需要真实值进行对齐,因此即使在测试阶段也可以灵活应用。
这种测试时集成方案能够在计算效率与预测质量之间提供灵活的权衡。通过选择合适的集成次数 N N N,用户可以在推理时间和深度估计精度之间找到最佳的平衡点。