简介:这是一份针对目标计数问题的扩散模型数据增强方案说明与可运行代码讲解,内容基于《Diffusion-based Data Augmentation for Object Counting Problems》进行复现。资源面向计算机视觉与深度学习从业者、研究者,以及希望借助扩散模型提升数据规模与多样性的算法学习者,尤其适合需要解决人群计数样本不足问题的场景。资源为单个docx文档,体积约24KB,以Python和PyTorch代码为主线,完整覆盖环境配置、预训练模型加载、密度图生成、条件损失与计数损失定义、最终合成人群图像等步骤。文档中演示了如何用高斯滤波将点标注转换为密度图,构造ControlNet训练流程,并通过对比合成与真实数据验证效果。目前已有63人学习,适合作为论文复现和技术调研的参考,能够帮助读者快速理解扩散模型在图像生成与计数任务中的结合方式,并获得可直接改用的实验思路。
1. 扩散模型数据增强,为什么目标计数先值得试
目标计数任务的常规数据增强,翻来覆去就是水平翻转、随机裁剪、颜色抖动这几招。人群计数的难点在于尺度变化、密集遮挡和背景纹理干扰,几何变换只改了坐标,并没有制造出「同样人群分布、不同外观」的新样本。扩散模型生成式增强的价值正在这里:以密度图为条件,重新生成对应的场景图像。密度图决定「哪里有人、有多少人」,扩散模型决定「这些人长什么样、背景是什么、遮挡关系如何」。下面这套路径从条件扩散的最小实现讲起,再落到计数网络训练、三组参数调优和样本筛选,适合正在做计数任务、想引入生成式增强但不想一上来就上大工程的人。
2. 从扩散模型原理到目标计数的适配:密度条件与流形空间
2.1 扩散模型原理是「走出流形再走回来」,数据增强等价于在流形邻域采样
扩散模型的核心过程可以压缩成两句话:前向过程逐步往图像上加噪,直到样本接近标准正态分布;反向过程学习一步步去噪,把纯噪声还原成图像。形式上,前向加噪是
x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon其中alpha_bar_t是第 t 步的累计噪声调度,epsilon是标准高斯噪声。反向过程用一个 UNet 预测每一步的噪声epsilon_theta(x_t, t),再用预测噪声反推上一步的图像。
用流形语言理解会更直观:真实图像分布集中在高维空间的一个低维流形上,前向加噪是把样本推离这个流形,反向去噪则是学习一条从噪声回到流形的路径。对数据增强来说,这意味着扩散模型不是「把一张图改一改」,而是在给定约束的前提下,从流形邻域里重新采样一张外观完全不同的图像。
这个性质和目标计数的契合点在于密度图。计数任务的标签不是类别,而是带空间分布的概率图。一张密度图对应的人头位置是确定的,但对应的人流外观、服装颜色、拍摄亮度可以千变万化。扩散模型在流形邻域采样时引入的随机性,恰好可以覆盖这些变化。相比之下,GAN 的单步生成是确定性的,难以在一次采样里同时撑起多个合理的视觉细节。
2.2 条件注入选型:通道拼接、FiLM 与潜在扩散模型控制分支
要让扩散模型按照密度图生成图像,必须把条件信息注入到去噪网络里。常见做法有三种。
第一种是通道拼接:把密度图插值成与输入图像相同尺寸,作为 UNet 的额外输入通道。图像 x 是 1 个通道,密度图 d 也取 1 个通道,拼接后变成 2 通道输入。这种实现方式保留了全部空间位置信息,代码最简单,适合小数据集和快速验证。缺点是没有专门的特征交互结构,条件信息在网络浅层混入后可能被后续卷积稀释。
第二种是 FiLM 调制:把密度图编码成一个条件向量,然后通过逐通道缩放和偏移来调节 UNet 特征图。FiLM 的参数效率高,但全局向量会丢失密度峰的具体位置信息,更适合「知道这一片大概有多少人」这种粗粒度条件。
第三种是潜在扩散模型中的控制分支:先用 VAE 把图像压缩到潜空间,在潜向量上做扩散,再用单独的编码器把密度图映射成空间控制特征,逐层注入 UNet。这是目前大图上效果最稳定的方案,但工程量和显存开销都更高。
三者的取舍可以用下面这张表概括:
| 注入方式 | 空间信息保留 | 实现成本 | 计数场景建议 |
|---|---|---|---|
| 通道拼接 | 完整 | 低 | 数据量小、先跑通流程时首选 |
| FiLM 调制 | 粗糙(全局向量) | 低 | 只关心总人数、不关心位置时使用 |
| 控制分支(潜在扩散) | 完整且分层 | 高 | 大分辨率、追求细节时使用 |
做目标计数时,密度峰的位置直接影响 MAE 指标,我一般先走通道拼接。它保留的空间信息最完整,调参简单,后面如果效果受限再切换成控制分支,网络结构改动也不大。
2.3 联合增强对象:为什么增强的是「图像 + 密度图」配对而不是单张图
传统几何增强里,翻转一张图必须同步翻转标签。在计数任务中这看起来合理,但本质上只是坐标变换,没有创造新的分布外样本:行人还是那些人,遮挡关系还是原来的遮挡关系。扩散增强的做法是训练一个条件生成模型p(x | d),用真实密度图 d 生成新的图像 x',配对关系保持为(x', d)。
有人会问:生成的图像里人多了一个或少了一个怎么办?这要回到计数任务的容错性上。密度图是空间概率分布,不是硬性的逐个标注,回归模型对局部峰值位置的细微偏移并不敏感;只要总体人数和大致分布保持一致,新样本就能参与训练。而分类任务里标签是一个离散类别,生成样本稍有歧义就会变成错误标注。这是计数任务先天适合扩散增强的重要原因。
实际实现时需要注意一个细节:生成样本只替换图像,不重新生成密度图。密度图由真实标注点加高斯核得到,如果对生成图像的密度图重新估计,会把计数模型的误差放进训练标签里,形成闭环放大。
3. 可运行的扩散增强与计数验证:最小 DDPM 条件生成代码
3.1 先用合成场景跑通数据接口:图像与密度图配对生成
为了把整套流程控制在可运行范围内,这里用合成场景代替真实数据集:目标用高斯椭圆模拟,密度图用点标注加高斯核生成。换成真实数据集时,只需要替换数据加载部分,后续网络和训练流程不需要改动。
import numpy as np def make_scene(size: int = 96, max_objects: int = 10, sigma: float = 2.0): x = np.arange(size) y = np.arange(size) xx, yy = np.meshgrid(x, y) # 背景:低频正弦 + 高频噪声,让模型学会区分前景目标与背景纹理 bg = 0.3 * np.sin(xx / 12.0) * np.cos(yy / 9.0) + 0.1 * np.random.randn(size, size) scene = bg.copy() density = np.zeros((size, size), dtype=np.float32) n = np.random.randint(3, max_objects) for _ in range(n): cx = np.random.randint(8, size - 8) cy = np.random.randint(8, size - 8) sx = np.random.uniform(2.0, 4.0) # 目标水平尺度 sy = np.random.uniform(2.0, 4.0) # 目标垂直尺度 theta = np.random.uniform(0, np.pi) a = np.square((xx - cx) * np.cos(theta) + (yy - cy) * np.sin(theta)) / (2 * sx * sx) b = np.square((xx - cx) * np.sin(theta) - (yy - cy) * np.cos(theta)) / (2 * sy * sy) obj = np.exp(-(a + b)) # 椭圆高斯目标 scene += obj * np.random.uniform(0.6, 1.0) density += np.exp(-(np.square(xx - cx) + np.square(yy - cy)) / (2 * sigma * sigma)) scene = np.clip(scene, 0, 1).astype(np.float32) density = np.clip(density, 0, None).astype(np.float32) return scene, density这个函数一次返回一对(图像, 密度图)。目标用带旋转角度的椭圆高斯模拟行人轮廓,密度图用固定sigma=2.0的高斯核渲染,保证每个目标在密度图上是一个独立峰值。背景用正弦低频分量加噪声,避免扩散模型偷懒把所有像素都预测成同一个纹理。
sigma参数值得专门说明:密度图中每个峰的宽度对应目标在图像中的实际大小。合成数据里可以固定,真实数据中如果场景有远近透视,应该按头部尺寸调整 sigma,否则一个目标会覆盖多个像素区域,计数网络学到的密度峰位置会系统性偏移。
3.2 加噪过程与条件 UNet:把密度图作为生成约束
接下来定义一个小规模 UNet。输入是拼接后的 2 通道张量(图像 + 密度图),时间步用正弦编码后通过 FiLM 层调制特征,输出是预测噪声。
import torch import torch.nn as nn import math class SinusoidalPosEmb(nn.Module): def __init__(self, dim: int = 128): super().__init__() self.dim = dim def forward(self, t): half = self.dim // 2 emb = math.log(10000) / (half - 1) emb = torch.exp(torch.arange(half, dtype=torch.float32, device=t.device) * -emb) emb = t[:, None].float() * emb[None, :] return torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) class SimpleUnet(nn.Module): def __init__(self, in_channels: int = 2, base: int = 32, time_dim: int = 128): super().__init__() self.time_mlp = nn.Sequential(SinusoidalPosEmb(time_dim), nn.Linear(time_dim, time_dim)) self.inc = nn.Conv2d(in_channels, base, 3, padding=1) self.down1 = nn.Sequential(nn.Conv2d(base, base * 2, 4, stride=2, padding=1), nn.GroupNorm(8, base * 2), nn.SiLU()) self.down2 = nn.Sequential(nn.Conv2d(base * 2, base * 4, 4, stride=2, padding=1), nn.GroupNorm(8, base * 4), nn.SiLU()) self.mid = nn.Sequential(nn.Conv2d(base * 4, base * 4, 3, padding=1), nn.GroupNorm(8, base * 4), nn.SiLU()) self.up1 = nn.ConvTranspose2d(base * 4, base * 2, 4, stride=2, padding=1) self.up2 = nn.ConvTranspose2d(base * 2, base, 4, stride=2, padding=1) self.out = nn.Conv2d(base, 1, 3, padding=1) self.to_gamma_beta = nn.Linear(time_dim, base * 2) def forward(self, x, t): tb = self.to_gamma_beta(self.time_mlp(t)) # [B, base*2] gamma = tb[:, : self.inc.out_channels].unsqueeze(-1).unsqueeze(-1) beta = tb[:, self.inc.out_channels:].unsqueeze(-1).unsqueeze(-1) h = self.inc(x) h = h * (1 + gamma) + beta # FiLM 调制 h1 = self.down1(h) h2 = self.down2(h1) h = self.mid(h2) h = self.up1(h) + h1 # 与下采样特征相加 h = self.up2(h) + h return self.out(h)噪声预测是 DDPM 的标准训练目标。模型不直接预测去噪后的图像,而是预测加入的噪声epsilon。原因在于预测x0时,模型容易把高频细节提前固定,后续去噪步无法修正,导致生成图出现整体偏色或结构漂移。预测噪声则保留了每一步的修正空间,训练也更稳定。
输入的x已经是 2 通道拼接张量,所以代码里in_channels=2。to_gamma_beta把时间步编码映射成每个通道的缩放和偏移系数,实现 FiLM 调制。这里没有额外处理密度图,因为密度图已经作为条件通道参与了所有卷积层的计算。
3.2.1 把时间步嵌入换成可学习嵌入
如果时间步数量固定且训练步数较多,也可以把正弦编码换成nn.Embedding(T, time_dim)。正弦编码的优点是不需要维护额外参数、对未见过的 t 也能平滑插值;可学习嵌入的表达能力更强,但在小数据上更容易过拟合时间步。合成数据场景里正弦编码足够。
3.3 训练循环与采样器:DDPM 最小实现
噪声调度采用 cosine schedule。它是 DDPM 的改进版,在接近 t=0 和 t=T 时加噪速度更平滑,避免线性 schedule 在训练后期出现过多纯噪声样本。
def cosine_beta_schedule(T=1000, s=0.008): steps = torch.arange(T + 1, dtype=torch.float32) / T alpha_bar = torch.cos((steps + s) / (1 + s) * math.pi / 2) ** 2 alpha_bar = alpha_bar / alpha_bar[0] beta = 1 - alpha_bar[1:] / alpha_bar[:-1] return torch.clip(beta, 1e-5, 0.02) T = 1000 betas = cosine_beta_schedule(T) alphas = 1 - betas alpha_bar = torch.cumprod(alphas, dim=0) def add_noise(x0, t): eps = torch.randn_like(x0) sqrt_ab = torch.sqrt(alpha_bar[t]).view(-1, 1, 1, 1) sqrt_1_ab = torch.sqrt(1 - alpha_bar[t]).view(-1, 1, 1, 1) return sqrt_ab * x0 + sqrt_1_ab * eps, eps def train_step(model, opt, x0, d0, drop_prob=0.1): b = x0.shape[0] t = torch.randint(0, T, (b,)) x_t, eps = add_noise(x0, t) cond = d0 if torch.rand(1).item() > drop_prob else torch.zeros_like(d0) x_in = torch.cat([x_t, cond], dim=1) pred = model(x_in, t) return nn.functional.mse_loss(pred, eps)drop_prob=0.1表示 10% 的概率把条件置零,这是 classifier-free guidance 的训练前提。没有这个置零操作,采样时就不能用 CFG 力度来控制生成图像对密度图的忠实程度。关于 CFG 的具体参数在第 4 章展开。
采样时从标准噪声出发,按 DDPM 反向迭代 200 步。步数可以从 1000 压缩到 200,但不要压到 50 以下,合成小图上 50 步会出现密度峰偏移(详细对比见 4.2)。
@torch.no_grad() def sample(model, cond, n_steps=200): model.eval() b, _, h, w = cond.shape x = torch.randn(b, 1, h, w) t_seq = torch.linspace(0, T - 1, n_steps, dtype=torch.long)[::-1] for i in t_seq: t_batch = torch.full((b,), i, dtype=torch.long) alpha = alphas[i].item() sqrt_ab = math.sqrt(alpha_bar[i].item()) sqrt_1_ab = math.sqrt(1 - alpha_bar[i].item()) pred = model(torch.cat([x, cond], dim=1), t_batch) x = (x - (1 - alpha) / sqrt_1_ab * pred) / math.sqrt(alpha) if i > 0: x = x + sqrt(1 - alpha) * torch.randn_like(x) model.train() return x反向迭代公式中,分子里的1 - alpha是当前步的噪声权重,除以sqrt(1 - alpha_bar)是为了把预测噪声从累计噪声尺度归一化到当前尺度。最后一步不加随机噪声,否则会引入额外的高频毛刺。
这套最小实现的超参数如下:
| 参数 | 值 | 说明 |
|---|---|---|
| 图像尺寸 | 96×96 | 合成数据默认值,真实数据按需调整 |
| 扩散步数 T | 1000 | 训练时随机采样 t |
| 采样步数 | 200 | 采样时按步长跨 t |
| base 通道数 | 32 | UNet 首层通道数,显存紧张可降到 16 |
| 时间嵌入维度 | 128 | 与 FiLM 输入维度一致 |
| 条件 dropout | 0.1 | CFG 训练必需 |
| 学习率 | 2e-4 | Adam 默认即可 |
| batch size | 16 | 96×96 输入在单卡上无压力 |
3.4 把增强样本接进计数回归器
生成样本最终要交给计数网络验证。这里用一个极简的卷积回归器:几层卷积加池化,输出 1/4 分辨率密度图,损失直接用密度图的 MSE。
class TinyCounter(nn.Module): def __init__(self): super().__init__() self.features = nn.Sequential( nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), ) self.head = nn.Conv2d(64, 1, 1) def forward(self, x): return self.head(self.features(x)) # 输出尺寸为输入的 1/4 pred_d = counter(x.unsqueeze(0)) loss = nn.functional.mse_loss(pred_d, d_target)输出密度图的总和就是预测人数,所以计数网络训练时不需要全连接层做分类,直接用回归 loss 即可。判断增强效果时,对比两组实验:一组只用真实数据训练,另一组混入扩散增强样本,比较验证集 MAE 和 MSE。收益不是固定值,要关注的是 MAE 是否下降、以及生成的样本是否让验证集上密度图总和分布更接近真实分布。
4. 三组必调参数与失败模式:从欠增强到模式坍塌
4.1 条件强度 scale:CFG 太高会重复目标
classifier-free guidance 在采样时同时计算有条件预测和无条件预测,并做外推:
eps_cfg = eps_uncond + scale * (eps_cond - eps_uncond)scale=0表示完全忽略条件密度图,采样结果就是随机场景。目标计数场景建议从1.0起步,最大不超过4.0。CFG 的 scale 和 GAN 里的truncation类似:值越大,样本越贴近条件,多样性越低。
| CFG scale | 生成效果 | 计数场景表现 |
|---|---|---|
| 0 | 场景随机,目标位置和密度图无关 | 增强无效,样本不配对 |
| 1.0 | 条件约束力适中,保留较多随机性 | 推荐起点,多样性较好 |
| 2.0 | 目标位置更贴合密度峰,纹理更清晰 | 效果最稳定 |
| 4.0 及以上 | 密度峰被过度强调,局部重复目标 | 生成图可能出现「一个峰一群人」的假密集区 |
判断 CFG 是否过高的简单办法:把生成的图像和条件密度图叠加显示,如果单个密度峰附近出现明显的多个目标轮廓粘连,说明 scale 偏大。此时计数网络会把一个峰认成多个目标,MAE 不降反升。
4.2 采样步数与多样化种子:推理开销与多样性的折中
训练时 T=1000,采样时并不需要走满 1000 步。合成场景下 200 步已经足够,真实数据上 50 步配 DDIM 会快很多,但步数减少会直接影响密度峰的位置精度。
| 采样步数 | 合成数据观察 | 建议 |
|---|---|---|
| 10 | 目标边缘模糊,密度峰向背景纹理漂移 | 不推荐 |
| 50 | 目标轮廓基本完整,局部峰位有偏移 | 快速实验可选 |
| 200 | 峰位稳定,纹理清晰 | 默认选择 |
采样步数对应的是反向去噪的离散间隔。步数减少意味着每一步需要在更长的噪声轨迹上做预测,对噪声估计误差的容忍度变低。真实数据上如果显存允许,我一般先用 50 步生成批量样本做初步筛查,筛过的样本再用 200 步重新生成,这样兼顾速度和质量。
还要注意随机种子的影响。同一个条件密度图,换一个种子生成的图像可能完全不同。计数增强需要的是多样性,所以每个真实密度图最好采样 3 到 5 张候选图,再经过第 5 章的筛选保留其中一部分。如果固定种子只生成一张,增强样本数量不够,计数网络很快会把生成分布也过拟合进去。
4.3 增强样本配比与过拟合:什么时候停
增强样本混入训练集的比例是另一个关键旋钮。常见做法是先按 30% 的混合比例把生成样本和真实样本放进同一个 DataLoader,观察验证集 MAE 曲线。如果验证损失下降但验证集上的总人数偏差持续偏大,说明生成样本的密度结构有问题,应当降低比例或回到 4.1 调低 CFG。
配比失衡的典型信号有三个。第一,验证集 MAE 下降但 MSE 上升,说明计数网络在部分样本上出现了大的离群误差。第二,在验证集上统计预测总人数的直方图,如果增强训练后的直方图出现双峰或整体右移,说明生成样本扰乱了计数网络的偏置。第三,对同一条件生成多张图,逐帧求像素级方差,方差趋近于零意味着采样模式已经退化。模式坍塌时只调整配比没有意义,要先改噪声采样步数或 CFG scale。
一个小建议:增强样本不要和真实样本做 1:1 混合,而是按 batch 内比例来配。比如每个 batch 取 70% 真实数据、30% 生成数据,这样既能控制比例,又能保证计数网络每个迭代都能看到真实数据。
5. 用一致性筛选锁定增强样本:密度总和与纹理阈值
5.1 一致性筛选:给增强样本设一道人数与纹理关卡
生成的图像不能直接全部加入训练集。合成场景下生成模型可能产生两类坏样本:一类是目标整体消失但密度图保留了若干个峰,另一种是图像过度平滑、没有可学习的高频纹理。两类样本都会让计数网络学到错误映射。筛选思路是用一个已经训练好的计数网络当裁判,结合图像纹理统计量做过滤。
def laplacian_var(x): kernel = torch.tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=x.dtype, device=x.device) kernel = kernel.view(1, 1, 3, 3) lap = torch.nn.functional.conv2d(x, kernel, padding=1) return lap.var(dim=(1, 2, 3)) def filter_synthetic(x_gen, d_cond, counter, err_thresh=0.08, lap_thresh=80): x_gen = x_gen.float() d_cond = d_cond.float() pred_sum = counter(x_gen).flatten(1).sum(1) cond_sum = d_cond.flatten(1).sum(1).clamp(min=1) rel_err = (pred_sum - cond_sum).abs() / cond_sum lap = laplacian_var(x_gen) keep = (rel_err < err_thresh) & (lap > lap_thresh) return x_gen[keep], d_cond[keep]这里的逻辑分两层。第一层是人数一致性:用当前计数模型预测生成样本的总人数,如果和条件密度图总人数的相对误差超过err_thresh=0.08,说明生成目标数量与密度图明显不匹配。第二层是纹理一致性:拉普拉斯方差低于阈值的图基本是平滑模糊的,放进训练集只会让计数网络把边缘纹理特征当成噪声忽略掉。
err_thresh和lap_thresh的取值与图像尺寸强相关。96×96 小图上 0.08 和 80 是合理的起调值;如果换成 512×512 的真实图片,拉普拉斯方差会大几个数量级,需要重新统计。更稳妥的做法是先看一批生成样本的 rel_err 和 lap 分布,用分位数而不是固定阈值来筛。常见做法是保留 rel_err 最小的前 80% 和 lap 最高的前 80%,两者的交集作为最终训练集。
这个「生成-筛选-重训」的循环可以迭代两到三轮:第一轮用 baseline counter 筛选得到的样本训练新 counter,第二轮用新 counter 重新筛选同批生成样本,被保留的样本会更聚焦。但迭代超过三轮后收益衰减明显,因为 counter 的误差分布已经被训练数据同化,筛选标准会失去区分度。最后判断增强有效与否,不看个别样图的视觉质量,而是看验证集 MAE、MSE 和预测总人数的直方图是否和真实分布对齐。把生成样本丢进验证集做密度直方图对比,是最不费时的健康检查。
本文还有配套的精品资源,点击获取