之前在图像超分项目里尝试用扩散模型增强细节时,最大的感受是:纹理变丰富了,但结果并不可控。模型会在原本平滑的区域脑补出不该存在的结构,导致人眼看似清晰,真实还原度却下滑。后来接触到 Uncertainty-Guided Latent Diffusion Models 这个方向,才意识到“给扩散模型引入不确定性引导,是缓解幻觉细节、提升重建忠实度的一条有效思路”。
本文会从超分任务的核心难点讲起,拆解 Uncertainty-Guided LDM 的方法逻辑,再给出一个可运行的简化 PyTorch 示例,覆盖环境准备、代码实现、训练验证、常见问题和工程建议。无论你是在做学术复现,还是想把扩散超分落地到实际业务,都值得往下看。
1. 背景与核心概念
1.1 超分辨率任务与“忠实”的细节
超分辨率(Super Resolution,SR)的目标是给定一张低分辨率图像 LR,恢复出对应的高分辨率图像 HR。传统退化过程可以抽象为:
LR = (HR * k) ↓ s + n其中k是模糊核,↓ s是下采样操作,n是噪声。超分模型本质上是在做这个退化过程的逆问题。
但逆问题往往不是唯一解。同一张 LR 图像,可以对应多张看起来合理的 HR 图像。比如一根头发丝在 LR 里可能只占 2-3 个像素,放大后到底朝哪个方向延伸、弧度多大,模型只能猜测。这就引出了超分领域非常核心的权衡:
- 保真度(Fidelity):重建结果是否在像素级接近真实 HR,常用 PSNR、SSIM 衡量。
- 感知质量(Perception):重建结果人眼看起来是否自然清晰,常用 LPIPS、FID 衡量。
传统 CNN 超分模型更偏向保真,但结果容易过度平滑,缺少真实纹理;生成模型尤其是扩散模型,能生成丰富细节,却可能在“猜”的过程中产生幻觉,导致重建内容与真实场景不符。
所谓Faithful Super Resolution,就是希望在保留生成模型细节表达能力的同时,让生成结果尽可能忠实于真实内容,不凭空虚构结构。这也是 Uncertainty-Guided Latent Diffusion Models 这一方向重点解决的问题。
1.2 潜在扩散模型为什么适合超分
扩散模型(Diffusion Model)通过逐步向数据添加噪声,再训练网络学习逆向去噪过程,从而学习到数据分布。基础流程可以理解成两个阶段:
- 正向过程:给干净图像逐步加噪,直到变成接近纯噪声。
- 反向过程:从纯噪声出发,逐步去噪,恢复出干净图像。
DDPM 是扩散模型的基础框架。它的理论清晰,生成质量高,但直接在像素空间执行去噪非常耗时。潜在扩散模型(Latent Diffusion Model,LDM)做了关键改进:先用 VAE 把图像压缩到低维潜在空间,再在潜在空间上执行扩散过程。
这样做的好处主要有两个:
- 计算量大幅下降,训练和推理都更快。
- 潜在空间的语义表达更紧凑,模型更容易学习高层次结构和纹理规律。
在超分任务中,LDM 通常把 LR 图像作为条件输入,引导模型在潜在空间生成对应的 HR 内容。这种方式能生成比传统 CNN 更真实的纹理,但如果条件利用不充分,模型就可能在细节上“自由发挥”。
1.3 Uncertainty-Guided 要解决的核心问题
Uncertainty-Guided LDM 的思路很直接:模型应该知道哪些区域的信息是可靠的,哪些区域是模糊不确定的,然后根据可靠程度决定生成策略。
图中边缘、纹理清晰的地方,LR 提供的约束强,扩散模型应该严格还原;而平滑区域、严重退化区域,LR 提供的约束弱,模型需要小心生成,避免编造无法由 LR 支撑的细节。
不确定性引导通常表现为一个额外分支或额外条件:
- 网络会预测每个像素/区域的不确定性值。
- 在训练时,不确定性用于调整损失权重,让模型更关注可靠区域。
- 在采样时,不确定性可以控制去噪强度或条件注入强度,减少不可靠区域的自由度过高。
本质上,这是给扩散模型加了一种“自我认知”能力,让生成过程更加可控。
2. 方法思路:不确定性引导扩散模型
2.1 不确定性来源:认知不确定性与偶然不确定性
在深度学习中,不确定性通常分为两类:
偶然不确定性(Aleatoric Uncertainty)由数据本身的内在随机性导致,例如 LR 图像中模糊、噪声和无纹理区域,无法通过增加训练数据消除。在超分任务中,这类不确定性主要来源于退化过程的随机性和下采样造成的信息丢失。
认知不确定性(Epistemic Uncertainty)由模型对未知区域的认知不足导致,可以通过更多数据、更大模型来降低。在超分任务中,当遇到训练分布之外的退化类型时,模型会产生较高的认知不确定性。
Uncertainty-Guided LDM 通常更关注偶然不确定性,因为超分中的信息丢失是客观存在的,模型要做的就是感知这种不可靠性,并据此调整生成策略。不过在工程实现中,也可以通过集成多个模型或 Dropout 等方式近似估计认知不确定性,两者可以结合使用。
2.2 不确定性图的构造方式
不确定性可以用不同的方式表达,常见设计包括:
1. 基于图像特征的经验估计
对输入 LR 图像计算局部方差、梯度幅度或边缘密度,得到一张空间分布的不确定性图。纹理稀疏、边缘模糊的区域不确定性高,结构清晰的区域不确定性低。这种方法的优点是轻量、可解释,适合工程落地。
2. 基于网络预测的显式估计
在去噪网络旁边增加一个不确定性预测头,输入 LR 和当前噪声特征,输出逐像素不确定性图。训练时使用不确定性加权损失,让网络自行学习哪些区域容易出错。这也是“Uncertainty-Guided”更严格意义上的做法。
3. 基于分布的隐式估计
在采样过程中对同一输入执行多次采样,计算多张输出之间的像素级方差。方差大说明模型对结果不确定,可以作为不确定性图。优点是几乎不用改模型结构,缺点是需要额外采样成本。
2.3 不确定性在训练与采样中的注入方式
得到不确定性图后,如何把它用起来,决定了整个方案的下限。常见注入方式包括:
1. 损失加权
这是最直接的方式。在训练时对 MSE 或 L1 损失按不确定性进行加权:
loss = mean( (pred_noise - noise)^2 * (1 + α * uncertainty) )其中α控制不确定性影响的强度。可靠区域权重更高,模型优先拟合这些区域;不可靠区域权重更低,避免模型在信息不足的区域强行过拟合。
2. 条件输入
把不确定性图作为额外的通道输入给去噪网络,与 LR 条件拼接。这样模型在每个去噪步骤中都能感知到各区域的可信度,从而调整注意力分配。
3. 生成强度控制
在采样阶段,如果不确定性高,可以适度降低噪声预测的置信度,或者对条件特征做软门控。例如通过不确定性图对条件特征加权:
cond_fused = lr_feature + gate * (1 - uncertainty) * generative_feature这样低不确定性区域更多依赖 LR,高不确定性区域才允许生成模型发挥。
3. 环境准备与项目结构
3.1 运行环境与依赖
本文示例使用 PyTorch 实现,核心依赖如下。版本需要根据实际环境调整,关键思想不受版本影响。
Python 3.8+ PyTorch 2.0+ torchvision Pillow numpy tqdm硬件方面,训练完整的 LDM 需要较高显存,建议至少 16GB。示例代码允许把图像裁剪为较小 patch 运行,显存不够时可以降到64x64或使用混合精度训练。
3.2 项目结构
一个清晰的项目结构能降低调试成本。本文示例目录如下:
uncertainty_guided_sr/ ├── config.py # 训练和采样参数 ├── diffusion.py # 噪声调度、加噪、采样 ├── uncertainty.py # 不确定性估计模块 ├── model.py # 条件去噪网络 ├── dataset.py # 数据集加载与退化 ├── train.py # 训练脚本 ├── infer.py # 推理脚本 └── sample_images/ # 推理输出目录为便于阅读,下面会把关键代码拆开讲解,并在代码前标注文件路径。
3.3 数据准备与退化设置
超分模型训练需要成对的 HR-LR 数据。如果使用通用数据集,读取 HR 图像后,可以动态生成 LR:
- 对 HR 图像做 4 倍双三次下采样得到 LR。
- 再把 LR 上采样回 HR 尺寸,作为条件输入。
- 按需添加高斯噪声或模糊,模拟真实退化。
这样做的好处是训练数据可以无限扩充,并且退化方式可控。
4. 核心代码拆解
4.1 噪声调度与前向加噪
扩散模型的第一步是定义加噪过程。这里使用经典的线性 beta 调度。
# 文件路径:diffusion.py import torch def linear_beta_schedule(timesteps, beta_start=1e-4, beta_end=0.02): return torch.linspace(beta_start, beta_end, timesteps) timesteps = 1000 betas = linear_beta_schedule(timesteps) alphas = 1.0 - betas alphas_cumprod = torch.cumprod(alphas, dim=0) sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)前向加噪过程定义为:
def q_sample(x_start, t, noise=None): # x_start: [B, C, H, W] # t: [B] if noise is None: noise = torch.randn_like(x_start) sqrt_alpha_bar = sqrt_alphas_cumprod[t].view(-1, 1, 1, 1) sqrt_one_minus = sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1) return sqrt_alpha_bar * x_start + sqrt_one_minus * noiset是每个样本对应的噪声步数索引,t越大,图像中保留的原始信息越少。模型要做的,就是根据带噪图像、步数和条件,预测出noise。
这种写法是 DDPM 的常规前向流程,也是后面训练数据准备的基础。
4.2 不确定性估计模块
先用轻量的局部统计量估计不确定性。这里以局部方差作为不确定性代理:局部方差越小,代表该区域越平滑,信息可靠性越低。
# 文件路径:uncertainty.py import torch import torch.nn.functional as F def estimate_uncertainty(img): """ 输入: img [B, C, H, W],值范围约 [0, 1] 输出: unc [B, 1, H, W],范围约 [0, 1] """ # 转为灰度图,减少通道计算 gray = img.mean(dim=1, keepdim=True) # [B, 1, H, W] # 3x3 局部均值 kernel = torch.ones(1, 1, 3, 3, device=img.device) / 9.0 mean = F.conv2d(gray, kernel, padding=1) mean_sq = F.conv2d(gray ** 2, kernel, padding=1) # 局部方差 var = torch.clamp(mean_sq - mean ** 2, min=0) # 归一化到 0~1 max_val = var.max(dim=2, keepdim=True)[0].max(dim=3, keepdim=True)[0] unc = var / (max_val + 1e-5) return unc在实际论文实现中,不确定性分支通常是从网络特征中学习出来的,效果会更准。这里用局部方差近似,主要目的是把引导机制跑通。你可以把它看成不确定性分支的替代品。
4.3 条件去噪模型
为了让模型能感知不确定性,我在去噪网络中增加了一个不确定性卷积分支,并把时间步嵌入加到隐藏层。
# 文件路径:model.py import torch import torch.nn as nn import torch.nn.functional as F def timestep_embedding(t, dim, max_period=10000): half = dim // 2 freqs = torch.exp( -torch.log(torch.tensor(max_period, dtype=torch.float32, device=t.device)) * torch.arange(half, dtype=torch.float32, device=t.device) / half ) args = t[:, None].float() * freqs[None, :] return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) class SimpleCondUNet(nn.Module): def __init__(self, in_ch=7, hidden_ch=64, time_dim=128): super().__init__() self.time_dim = time_dim self.time_embed = nn.Sequential( nn.Linear(time_dim, time_dim * 2), nn.SiLU(), nn.Linear(time_dim * 2, time_dim * 2), nn.SiLU(), ) self.time_proj = nn.Linear(time_dim * 2, hidden_ch) self.unc_conv = nn.Conv2d(1, hidden_ch, 3, padding=1) self.conv1 = nn.Conv2d(in_ch, hidden_ch, 3, padding=1) self.conv2 = nn.Conv2d(hidden_ch, hidden_ch, 3, padding=1) self.conv3 = nn.Conv2d(hidden_ch, hidden_ch, 3, padding=1) self.out = nn.Conv2d(hidden_ch, 3, 3, padding=1) def forward(self, x_noisy, t, lr_cond, unc_map): # x_noisy: [B, 3, H, W] # lr_cond: [B, 3, H, W] # unc_map: [B, 1, H, W] # 时间步嵌入 emb = timestep_embedding(t, self.time_dim) emb = self.time_embed(emb) emb = self.time_proj(emb)[:, :, None, None] # [B, hidden, 1, 1] # 不确定性条件特征 unc_feat = self.unc_conv(unc_map) # [B, hidden, H, W] # 拼接输入:噪声图 + LR + 不确定性图 h = torch.cat([x_noisy, lr_cond, unc_map], dim=1) h = F.silu(self.conv1(h)) h = h + emb + unc_feat h = F.silu(self.conv2(h)) h = F.silu(self.conv3(h)) return self.out(h)这个模型是简化版,用于演示核心机制。真实 LDM 中的 UNet 会有下采样、上采样、Attention 和 skip connection,如果你复现论文,需要把这里的SimpleCondUNet替换成完整的 UNet 或使用开源实现。
输入通道为 7,分别对应:
- 3 通道:带噪 HR 图像
- 3 通道:上采样后的 LR 条件
- 1 通道:不确定性图
时间步嵌入让网络知道当前去噪阶段,不确定性分支让网络知道每块区域的可靠度。
4.4 不确定性加权损失
训练损失不能只做简单的像素级 MSE。我们希望模型优先拟合信任区域,弱化不可靠区域的影响。
def uncertainty_weighted_loss(pred_noise, noise, unc_map, alpha=1.0): # pred_noise: [B, 3, H, W] # noise: [B, 3, H, W] # unc_map: [B, 1, H, W] mse = F.mse_loss(pred_noise, noise, reduction="none") # [B, 3, H, W] mse = mse.mean(dim=1) # 转为逐像素损失 [B, H, W] # 高不确定性区域权重更低 weight = 1.0 + alpha * (1.0 - unc_map.squeeze(1)) # [B, H, W] loss = (mse * weight).mean() return loss这个设计表达了一个直觉:LR 中结构清晰的区域,模型预测的噪声应该被严格监督;而模糊、不确定的区域,模型“猜错”的代价降低,避免强迫模型把不可靠信息当成确定信息来学。
4.5 采样推理
训练完成后,从一个随机噪声图出发,逐步去噪得到 SR 结果。采样时同样把 LR 条件和不确定性图送入模型。
@torch.no_grad() def sample(model, lr_up, unc_map, device, timesteps=1000): model.eval() b = lr_up.size(0) x = torch.randn_like(lr_up).to(device) for i in reversed(range(timesteps)): t = torch.full((b,), i, device=device, dtype=torch.long) pred_noise = model(x, t, lr_up, unc_map) alpha = alphas[i] alpha_bar = alphas_cumprod[i] sqrt_one_minus = sqrt_one_minus_alphas_cumprod[i] x = (1.0 / torch.sqrt(alpha)) * ( x - (1.0 - alpha) / sqrt_one_minus * pred_noise ) if i > 0: beta = betas[i] x = x + torch.sqrt(beta) * torch.randn_like(x) return x采样时每一步都参考同一张不确定性图,所以在结构清晰区域,模型受 LR 条件的约束更稳定;在不确定区域,模型能通过生成补充细节,但不会无限发挥。
5. 训练与验证完整流程
5.1 数据集加载
下面是一个简单的 SR 数据集实现。它读取 HR 图片,随机裁剪 patch,然后动态生成 LR。
# 文件路径:dataset.py import os import torch from torch.utils.data import Dataset from torchvision import transforms from PIL import Image class SRDataset(Dataset): def __init__(self, root, patch_size=128): self.root = root self.paths = [ os.path.join(root, f) for f in os.listdir(root) if f.lower().endswith((".png", ".jpg", ".jpeg")) ] self.patch_size = patch_size self.to_tensor = transforms.ToTensor() def __len__(self): return len(self.paths) def __getitem__(self, idx): img = Image.open(self.paths[idx]).convert("RGB") # 随机裁剪 w, h = img.size if w >= self.patch_size and h >= self.patch_size: i = torch.randint(0, h - self.patch_size + 1, ()) j = torch.randint(0, w - self.patch_size + 1, ()) img = img.crop((j, i, j + self.patch_size, i + self.patch_size)) else: img = img.resize((self.patch_size, self.patch_size), Image.BICUBIC) hr = self.to_tensor(img) # [3, H, W] return {"hr": hr}使用时要保证图片路径存在。真实项目可以换成自定义列表,读取 txt/csv 中的路径。
5.2 训练循环
训练脚本整合了前向加噪、不确定性估计、模型前向和损失计算。
# 文件路径:train.py import torch import torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import transforms from diffusion import q_sample, betas, alphas_cumprod from uncertainty import estimate_uncertainty from model import SimpleCondUNet from dataset import SRDataset device = "cuda" if torch.cuda.is_available() else "cpu" timesteps = 1000 epochs = 20 batch_size = 8 lr = 1e-4 patch_size = 128 model = SimpleCondUNet().to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=lr) dataset = SRDataset(root="./data/train", patch_size=patch_size) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4) model.train() for epoch in range(epochs): total_loss = 0.0 for batch in dataloader: hr = batch["hr"].to(device) # [B, 3, H, W] # 生成 LR 并上采样回 HR 尺寸 lr = F.interpolate(hr, scale_factor=0.25, mode="bicubic", align_corners=False) lr_up = F.interpolate(lr, size=(hr.size(2), hr.size(3)), mode="bicubic", align_corners=False) # 估计不确定性 unc = estimate_uncertainty(lr_up) # 随机时间步 t = torch.randint(0, timesteps, (hr.size(0),), device=device).long() noise = torch.randn_like(hr) noisy_hr = q_sample(hr, t, noise) pred_noise = model(noisy_hr, t, lr_up, unc) loss = F.mse_loss(pred_noise, noise, reduction="none") loss = loss.mean(dim=1) # [B, H, W] weight = 1.0 + 1.0 * (1.0 - unc.squeeze(1)) loss = (loss * weight).mean() optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f"Epoch {epoch + 1}/{epochs} - loss: {total_loss / len(dataloader):.6f}")如果显存不足,可以降低patch_size或batch_size。如果模型收敛慢,可以检查 lr 是否合适,以及不确定性图是否被过度归一化。
5.3 推理脚本
推理脚本加载 HR 图片,生成 LR,再调用采样函数得到 SR 结果,和真实 HR 对比。
# 文件路径:infer.py import torch import torch.nn.functional as F from torchvision import transforms from PIL import Image from diffusion import betas, alphas, alphas_cumprod, sqrt_one_minus_alphas_cumprod from uncertainty import estimate_uncertainty from model import SimpleCondUNet device = "cuda" if torch.cuda.is_available() else "cpu" timesteps = 1000 model = SimpleCondUNet().to(device) model.load_state_dict(torch.load("model.pth", map_location=device)) to_tensor = transforms.ToTensor() to_pil = transforms.ToPILImage() img = Image.open("./test.png").convert("RGB") hr = to_tensor(img).unsqueeze(0).to(device) lr = F.interpolate(hr, scale_factor=0.25, mode="bicubic", align_corners=False) lr_up = F.interpolate(lr, size=(hr.size(2), hr.size(3)), mode="bicubic", align_corners=False) unc = estimate_uncertainty(lr_up) sr = sample(model, lr_up, unc, device, timesteps) sr = torch.clamp(sr, 0, 1) to_pil(sr.squeeze(0)).save("./sample_images/result.png")运行后可以肉眼对比result.png与真实 HR 之间的差异。也可以继续计算 PSNR、SSIM、LPIPS 等指标。
6. 常见问题与排查思路
扩散模型训练本身就有不少坑,加入不确定性引导后又会多出一些新问题。下面整理几个高频场景:
| 问题现象 | 常见原因 | 解决思路 |
|---|---|---|
| 训练 loss 不下降 | 学习率过大或过小;不确定性图归一化异常 | 调小学习率,打印不确定性图的数值分布,确认范围在 0~1 |
| 生成结果过于模糊 | 条件注入太强,模型只依赖 LR;去噪步数不足 | 降低条件权重,增加步数,检查不确定性加权是否过度压制生成 |
| 生成结果出现伪细节 | 不确定性引导失效,模型在平滑区域过度发挥 | 提高不可靠区域损失压制强度,或改用网络预测的不确定性分支 |
| 采样结果噪声明显 | 推理时未加上噪声调度或最后一步处理错误 | 检查sample函数中beta是否使用正确,最后一轮不要加噪声 |
| GPU 显存不足 | patch 太大或 batch 太大 | 降低patch_size到 64,开启 AMP 混合精度 |
| 不确定性图变成全 0 | 局部方差计算后归一化除零 | 检查var.max(),添加1e-5防止除零 |
| 结果与 LR 内容不对齐 | 条件信息不足或模型没有多尺度结构 | 使用真实 UNet + Attention,增加 LR 多尺度特征拼接 |
排查时建议先做最小实验:固定 schedule,去掉不确定性加权,跑 1000 步确认基础流程能收敛,再加入不确定性分支做对比。
7. 最佳实践与工程建议
7.1 数据与退化设置
训练扩散超分模型的退化设置非常关键。如果只用双三次下采样,模型面对真实退化时会退化明显。建议采用随机退化组合:
- 随机模糊核(高斯、运动模糊)
- 随机噪声(高斯噪声、JPEG 压缩噪声)
- 随机下采样倍率
退化类型越接近线上数据,不确定性估计越有实际意义,模型的忠实度也越高。
7.2 训练稳定性
扩散模型的训练稳定性与噪声调度、批大小、学习率强相关。经验上,建议使用AdamW,初始学习率从1e-4到3e-4,配合梯度裁剪max_grad_norm=1.0。
如果显存允许,尽量增大 batch size;如果 batch 受限,可以降低 patch size,但不要低于 64。否则批量归一化和卷积的感受野都会受到明显影响。
7.3 评估指标
超分模型不能只看 PSNR。对 Uncertainty-Guided LDM 这类生成式模型,建议同时关注:
- PSNR、SSIM:衡量保真度。
- LPIPS:衡量感知相似度。
- FID:衡量生成分布与真实分布的接近程度。
- 忠实度专项指标:例如边缘误差、结构相似度、区域级一致性。
更严格的做法是人工评测,重点看平坦区域是否出现伪造纹理、边缘是否清晰且贴合 LR 结构。
7.4 部署与加速
扩散模型推理速度慢是落地的主要瓶颈。可以从几个方向优化:
- 使用预训练 VAE 在潜在空间操作,减少计算量。
- 采用 DDIM 采样或 DPM-Solver,减少推理步数。
- 蒸馏为一步或少步模型,例如 LCM、Turbo 等少步采样方案。
- 使用 TensorRT 或 ONNX Runtime 部署,配合 FP16。
需要注意的是:减少采样步数会对生成忠实度产生影响,部署前必须做充分验证,尤其是高不确定区域的细节表现。
8. 下一步学习路线
如果你看完本文后,想继续深入 Uncertainty-Guided LDM 这个方向,建议按下面的顺序推进:
- 先跑通本文的简化版代码,理解 DDPM 加噪、去噪、采样全过程。
- 阅读 DDPM / DDIM 原论文,搞清楚每一步的公式推导。
- 用开源 Stable Diffusion 的 VAE 和 UNet 替换简化模型,把像素空间训练改成潜在空间训练。
- 在潜在空间加入不确定性分支,对比有/无不确定性引导的生成效果。
- 在真实业务数据上测试模糊、压缩、混合退化场景,统计不同区域的不确定性与重建误差关系。
- 尝试用不确定性图控制采样的 guidance scale,探索更细粒度的生成控制策略。
做扩散超分落地时,最需要警惕的就是“看起来清晰但内容不可靠”的假象。把不确定性引导加进去之后,模型并不是变得更聪明,而是变得更懂得“说不知道”。这种能力,恰恰是忠实超分最稀缺的。