ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distilla代码梳理

一、训练模块

1. 准备数据加载器

1.1 调用NeRFDataset中的dataloader对数据进行加载


① 返回加载器
② loader中包含的数据信息；图像高度、宽度、光线原点、光线方向、方向向量、姿态矩阵和投影矩阵等

2. 训练

2.1 调用train函数开始训练

2.2 train_one_epoch用于处理一个完整的训练周期(epoch)内的数据


①每隔16步更新密度网格和位字段。
②调用train_step获取预测的RGB图像、预测的深度图、损失值loss、伪损失值、图像潜在表示和渲染模式

如何获得的损失loss？

train_one_epoch中调用train_step(对每个数据批次进行具体计算)，然后在train_step中调用扩散模型文件(sd.py)中的train_step(用于扩散模型的单步训练步骤，主要负责生成潜在表示并计算损失值)

①添加噪声，并调用预训练的扩散模型 (冻结的StableDiffusion) 预测噪声noise_pred(基于文本)
②CFG通过混合无条件/有条件预测(加权融合)，增强文本控制能力
③调用微调后的 LoraUnet 进行噪声预测noise_pred_q(基于姿态和着色)
④在特定配置下开启预测噪声

①扩散模型（基于文本）和 微调模型（基于姿态）对同一潜变量的噪声估计差作为梯度
②通过自定义的自动求导操作计算损失值（绕过自动微分以实现精确控制，常见于扩散模型）

根据loss通过反向传播更新Nerf模型参数，优化粒子

通过缓冲区或者直接加载数据。

• 直接使用 data 会导致模型只学习当前批次的数据分布，容易过拟合。
• 通过缓冲区采样，模型可以学习到更广泛的数据分布，提升泛化能力

①对潜变量进行加噪处理，并调用微调后的 LoraUnet 进行噪声预测model_output(基于姿态和着色)
②通过均方误差(mse_loss)计算损失loss_unet
③反向传播更新LoraUnet的参数

每训练一轮(epoch)，就对模型进行评估和保存一次

训练一轮(epoch)，就根据着色方式(shading)进行图像的生成并保存

补充：

• LoraUnet.forward是LoRA模块的入口函数，它的职责是注入额外条件（如 c, shading）。
• 但它仍然需要支持文本引导，因为：
  • LoRA训练阶段依赖文本引导的噪声预测。
  • self.unet是标准扩散模型，必须支持encoder_hidden_states。
• LoraUnet.forward中虽然没有直接使用encoder_hidden_states，但在return self.unet(...)中仍然传入它，是为了保留文本引导能力，确保在NeRF训练阶段可以使用标准扩散模型进行监督。

lora_unet中标准的扩散模型的前向传播

二、变分分数蒸馏的更新规则

① Lora参数更新 ②Nerf粒子更新

1.Lora参数更新

$$\underset{\varphi }{\min }\sum _{i=1}^n{\mathbb{E}}_{t\sim \mathcal{U}(0,1),\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}),c\sim p(c)}\left[\Vert {\mathbf{ϵ}}_{\varphi }({\alpha }_t\mathbf{g}({\theta }^{(i)},c)+{\sigma }_t\mathbf{ϵ},t,c,y)-\mathbf{ϵ}{\Vert }_2^2\right]$$
$\mathbf{g}({\theta }^{(i)},c)$：第 $i$ 个粒子 $\theta$ 在相机位姿 $c$ 下渲染成的2D图像
${\alpha }_t$：累积噪声比例系数，由DDIM调度器决定
${\alpha }_t\mathbf{g}({\theta }^{(i)},c)+{\sigma }_t\mathbf{ϵ}$：第t步的带噪图像
$t$：扩散时间步t
$c$：相机位姿c
$y$：文本嵌入向量
${\mathbf{ϵ}}_{\varphi }(...)$：在参数 $\varphi$ 下Lora预测的噪声
$\mathbf{ϵ}$：$\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$标准高斯噪声
$\Vert {\mathbf{ϵ}}_{\varphi }(...)-\mathbf{ϵ}{\Vert }_2^2$：预测噪声与真实噪声的均方误差(MSE)
$\mathbb{E}$：对随机采样部分求期望
$\underset{\varphi }{\min }$：最小化损失函数

1. 2D图像渲染+潜在表示
   
   rays_o、rays_d：光线的起点和方向
   mvp：模型-视图-投影矩阵（Model-View-Projection），将3D物体坐标变换为屏幕坐标，供后续光栅化使用
   H、W：渲染图像的高度和宽度
   staged：是否分阶段返回结果（逐批次处理）
   light_d：定向光源方向，用于计算漫反射
   perturb：是否在光线采样时添加随机扰动
   bg_color：背景色
   ambient_ratio：环境光强度比例
   shading：着色模式
   binarize：是否将输出二值化（用于深度图）
   
   返回的结果中包含深度图depth、渲染图image、权重累加和weights_sum(光线被场景遮挡的累计权重，几何信息/透明度)
   
   当潜空间表示时，预测的图像pred_rgb拼接image(RGB颜色信息)和weights_sum(几何信息)生成4通道图像，否则使用image 数据，生成标准3通道RGB图像
   
   获取图像的潜在表示
   
   ①直接下采样(as_latent)最快速但丢失细节信息，可能导致生成结果缺乏精细结构​（如纹理、边缘清晰度）
   ②原生输入(opt.latent)完全保留信息但计算负担大
   ③通过变分自编码器(VAE)将RGB图像压缩为结构化潜向量​，在效率和保真度间取得平衡

2. 加噪+噪声预测
   
   生成的高斯噪声noise：$\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$标准高斯噪声
   加噪latents_noisy：${\alpha }_t\mathbf{g}({\theta }^{(i)},c)+{\sigma }_t\mathbf{ϵ}$
   Lora预测的噪声model_output：${\mathbf{ϵ}}_{\varphi }({\alpha }_t\mathbf{g}({\theta }^{(i)},c)+{\sigma }_t\mathbf{ϵ},t,c,y)$
   
   Lora在调用Unet预测噪声时，会调用Lora自己封装的Unet(LoraUnet)，在进行前向传播时会用到y(textemb)作为引导的。

3. 计算损失+反向传播更新参数
   
   在DDIM中，velocity是潜在空间中的一种特殊表示，它结合了原始图像和噪声的信息。
   get_velocity()并非生成新噪声，而是根据已有的latents_clean和 noise 计算等价的速度向量$v_t$
   计算损失(MSE)：$\sum _{i=1}^n{\mathbb{E}}_{t\sim \mathcal{U}(0,1),\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}),c\sim p(c)}\left[\Vert {\mathbf{ϵ}}_{\varphi }(...)-\mathbf{ϵ}{\Vert }_2^2\right]$
   反向传播更新参数(backward)：$\underset{\varphi }{\min }$，最小化损失值的过程

2.Nerf粒子更新

$${\nabla }_{\theta }{\mathcal{L}}_{\mathrm{VSD}}(\theta )\triangleq {\mathbb{E}}_{t,\mathbf{ϵ},c}\left[\omega (t)\left({\mathbf{ϵ}}_{\text{pretrain }}\left({\mathbf{x}}_t,t,y^c\right)-{\mathbf{ϵ}}_{\varphi }\left({\mathbf{x}}_t,t,c,y\right)\right)\frac{\partial \mathbf{g}(\theta ,c)}{\partial \theta }\right]其中{\mathbf{x}}_t={\alpha }_t\mathbf{g}(\theta ,c)+{\sigma }_t\mathbf{ϵ}$$
${\nabla }_{\theta }{\mathcal{L}}_{\mathrm{VSD}}(\theta )$：关于粒子参数 $\theta$ 的 $\mathrm{VSD}$ 损失函数的梯度

${\mathbf{x}}_t$：t时刻的带噪图像

$y^c$：带有方向的文本嵌入向量

$y$：文本嵌入向量

${\mathbf{ϵ}}_{\text{pretrain }}\left({\mathbf{x}}_t,t,y^c\right)$：预训练扩散模型预测的噪声

${\mathbf{ϵ}}_{\varphi }\left({\mathbf{x}}_t,t,c,y\right)$：在参数$\varphi$下Lora预测的噪声

$\omega (t)$：关于时间步t的权重函数，用于平衡不同噪声水平下梯度信号的重要性

$\frac{\partial \mathbf{g}(\theta ,c)}{\partial \theta }$：渲染图像$\mathbf{g}(\theta ,c)$关于Nerf模型参数 $\theta$ 的梯度,它是连接扩散模型预测与NeRF参数的桥梁。前面计算出的噪声差异作为损失信号，通过这个梯度反向传播回NeRF的参数θ。这告诉NeRF：为了减少两个噪声预测的差异（即让渲染的图像更符合通用扩散模型在条件c下的预期），应该如何改变自己的参数θ 。

1. 预训练扩散模型预测噪声
   
   
   ①text_embeddings是函数get_text_embeds返回的结果，她包含无文本(空)嵌入向量(uncond_embeddings)和有文本嵌入向量(text_embeddings)。unet同时学习无条件和有条件预测。
   ②noise_pred_uncond：无条件下预测的噪声，noise_pred_text：文本条件下预测的噪声，guidance_scale：CFG文本引导超参，强化文本条件对生成的引导作用。

noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
等于
noise_pred = (1 - guidance_scale ) * noise_pred_uncond + guidance_scale  * noise_pred_text
# 无条件和文本引导噪声的权重

2. Lora预测噪声，并开启v-prediction速度场
   v-prediction速度预测，是扩散模型中一种替代噪声预测的参数化方式。在传统的扩散模型中，模型通常预测添加到数据中的噪声($ϵ$)。但v_pred是一种替代参数化方法，模型预测的是速度项($v$)，它与噪声和潜在变量有线性关系。
   在扩散模型中，速度场的定义为v-prediction ：
   $$v_t=\sqrt{{\bar{\alpha }}_t}\cdot ϵ-\sqrt{1-{\bar{\alpha }}_t}\cdot x_0$$
   从$v$恢复$ϵ$:
   $$\begin{array}{c}{x}_t=\sqrt{{\bar{\alpha }}_t}\cdot x_0+\sqrt{1-{\bar{\alpha }}_t}\cdot ϵ\\ x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}\cdot x_t-\frac{\sqrt{1-{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_t}}\cdot ϵ\\ v_t=\sqrt{{\bar{\alpha }}_t}\cdot ϵ-\sqrt{1-{\bar{\alpha }}_t}\cdot \left(\frac{1}{\sqrt{{\bar{\alpha }}_t}}\cdot x_t-\frac{\sqrt{1-{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_t}}\cdot ϵ\right)\\ v_t=\sqrt{{\bar{\alpha }}_t}\cdot ϵ-\frac{\sqrt{1-{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_t}}\cdot x_t+\frac{(1-{\bar{\alpha }}_t)}{\sqrt{{\bar{\alpha }}_t}}\cdot ϵ\\ v_t=\left(\sqrt{{\bar{\alpha }}_t}+\frac{1-{\bar{\alpha }}_t}{\sqrt{{\bar{\alpha }}_t}}\right)\cdot ϵ-\frac{\sqrt{1-{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_t}}\cdot x_t\\ v_t=\left(\frac{1}{\sqrt{{\bar{\alpha }}_t}}\right)\cdot ϵ-\frac{\sqrt{1-{\bar{\alpha }}_t}}{\sqrt{{\bar{\alpha }}_t}}\cdot x_t\\ \Rightarrow ϵ=\sqrt{{\bar{\alpha }}_t}\cdot v_t+\sqrt{1-{\bar{\alpha }}_t}\cdot x_t\end{array}$$
   
   ①先通过Lora预测噪声noise_pred_q，也就是公式中的$v_t$
   ②$\sqrt{{\bar{\alpha }}_t}\cdot v$用模型预测的速度场还原噪声，t↑ ⇒ 趋于 0
   ③$\sqrt{1-{\bar{\alpha }}_t}\cdot x_t$用现有带噪图像自身补上噪声，t↑ ⇒ 趋于 1
   ④将 Lora 预测的$v_t$转换成$ϵ$ ，用于训练损失对比

3. 梯度计算+反向传播
   
   ①$\omega (t)$：关于时间步t的权重函数，高噪声步（t→T）梯度权重小，低噪声步（t→0）权重大，突出细节优化阶段用于平衡不同噪声水平下梯度信号的重要性
   ②$\omega (t)\left({\mathbf{ϵ}}_{\text{pretrain }}(...)-{\mathbf{ϵ}}_{\varphi }(...)\right)$ 计算指导生成的梯度
   $\frac{\partial \mathbf{g}(\theta ,c)}{\partial \theta }$用 grad作为 latents ($g(\theta ,c)$) 的梯度，强制地指导网络往想要的方向优化
   反向传播更新NeRF参数

① ② ③ ④ ⑤ ⑥ ⑦ ⑧ ⑨ ⑩