DSNT可微空间到数值变换：从热力图到坐标的端到端关键点检测

1. 从坐标回归的痛点说起做视觉关键点检测的朋友大概率都经历过这样的场景模型训练得好好的loss 曲线也漂亮但一到推理阶段把热力图转成坐标精度就莫名其妙掉了一截。更让人抓狂的是这个转换过程不可导梯度传不回去端到端训练直接断链。我最早做人脸关键点的时候用的是经典的 argmax 取峰值方案训练阶段模型优化的是热力图上的高斯分布推理阶段却用 argmax 去取整数坐标两者之间存在一个天然的鸿沟——训练时模型根本不知道你最后要用 argmax 取点它只管把热力图做好看至于取出来的坐标准不准跟它的优化目标没有直接关系。这个问题困扰了相当长一段时间。后来接触到 DSNTDifferentiable Spatial to Numerical Transform才算是找到了一个比较优雅的解法。DSNT 的核心思路非常直接设计一个可微的变换层把空间上的热力图spatial映射成数值坐标numerical而且整个过程保持可导让梯度可以从坐标的 loss 一路回传到热力图的每一个像素上。这样一来模型在训练时就能感知到最终的坐标误差优化目标从热力图像不像变成了坐标准不准端到端的一致性一下子就建立起来了。这篇文章适合谁看如果你正在做人体姿态估计、人脸对齐、手部关键点、医学图像定位这类需要从热力图回归坐标的任务并且对精度有比较高的要求那 DSNT 值得你花时间研究。如果你只是刚入门关键点检测还没遇到 argmax 精度损失的问题也可以先了解一下这个思路等踩到坑了再回来翻。全文我会从设计动机、数学原理、代码实现、训练技巧、常见问题几个维度展开尽量把每个为什么都讲清楚让你看完能直接上手改自己的模型。2. DSNT 到底解决了什么问题2.1 传统热力图方法的精度瓶颈先说说传统方案为什么会有精度问题。主流的关键点检测方法大致分两派一派是直接回归坐标direct regression网络最后接几个全连接层直接输出 x、y 数值另一派是热力图heatmap-based网络输出 K 张热力图每张对应一个关键点然后从热力图上取峰值位置作为坐标。直接回归的问题在于坐标是一个高度非线性的映射网络要直接从图像特征回归出精确数值学习难度大而且缺乏空间上的可解释性训练不稳定。热力图方法则把回归问题转化成了逐像素的分类/密度估计问题空间结构保留得好训练更稳定精度通常也更高。所以现在主流方案基本都是热力图路线。但热力图路线有个绕不开的环节从热力图到坐标的转换。最常见的做法是 argmax取响应最大的那个像素位置。这个操作有两个致命问题。第一argmax 的输出是整数只能精确到像素级别亚像素精度直接丢失。第二argmax 不可导梯度无法回传训练和推理之间存在断层。有人会说那我用 soft-argmax 不就行了soft-argmax 确实可导它用 softmax 把热力图归一化成概率分布然后对坐标做加权平均。这个思路已经比 argmax 好很多了但 soft-argmax 有个隐患softmax 的归一化会改变热力图的空间分布形状尤其是当热力图比较尖锐或者存在多个峰值时softmax 会把概率质量重新分配导致加权平均出来的坐标产生系统性偏移。而且 softmax 对温度参数敏感温度设不好要么分布太平导致坐标被拉向中心要么分布太尖导致梯度消失。2.2 DSNT 的核心设计思路DSNT 的做法可以拆成两步。第一步对热力图做归一化让它变成一个合法的概率分布所有元素非负且和为 1。第二步用两个固定的坐标期望矩阵分别和归一化后的热力图做点积得到 x 和 y 的期望值也就是最终坐标。这里的关键在于归一化用的是 L2 归一化或者叫归一化到和为 1 的任意可微方式而不是 softmax。L2 归一化不会像 softmax 那样引入指数运算带来的分布形变它只是把热力图整体缩放到和为 1保留了原始的空间相对关系。然后坐标期望的计算是线性的梯度可以非常干净地回传。用一句话概括 DSNT 的哲学与其让网络去拟合一个好看的热力图再事后取点不如让网络直接优化坐标的期望热力图只是中间产物。这个视角的转变是 DSNT 区别于其他方案的根本所在。2.3 和其他方案的对比为了让你更直观地理解 DSNT 的定位我整理了一个对比表方案可导性亚像素精度分布形变实现复杂度训练稳定性argmax不可导无无极低训练推理不一致soft-argmax可导有有softmax 导致低中等直接回归可导有不适用低较差DSNT可导有无L2 归一化低好从表里能看出来DSNT 在可导性、亚像素精度、分布保真度这几个维度上都做到了比较好的平衡。实现复杂度也不高核心代码就几十行这也是它比较受欢迎的原因之一。3. DSNT 的数学原理拆解3.1 归一化把热力图变成概率分布假设网络输出的热力图是 H尺寸为 m×n每个元素 H[i][j] 表示该位置是目标点的响应强度。DSNT 第一步是把它归一化成一个概率分布 P满足 P[i][j] ≥ 0 且 ΣP[i][j] 1。最直接的做法是除以总和P[i][j] H[i][j] / Σ H[a][b]但这里有个坑如果 H 里存在负值直接除会导致 P 出现负值不再是合法分布。所以实践中通常先对 H 做一次非负变换比如取平方或者取绝对值再归一化。取平方还有个额外好处它会放大高响应区域、抑制低响应区域让分布更集中坐标期望更接近真实峰值。我个人的经验是如果热力图是用 ReLU 或者 sigmoid 激活输出的本身已经非负那直接除以总和就行。如果是线性输出建议先平方再归一化。这个细节看起来小但对最终精度有影响后面讲实操的时候会再展开。3.2 坐标期望用固定矩阵做点积归一化得到 P 之后接下来计算坐标期望。这里需要两个固定的矩阵X 矩阵和 Y 矩阵尺寸都是 m×n。X 矩阵的构造方式是X[i][j] j / (n-1)也就是每个位置的列索引归一化到 [0, 1]。Y 矩阵类似Y[i][j] i / (m-1)行索引归一化到 [0, 1]。然后坐标期望就是x_hat Σ P[i][j] * X[i][j] y_hat Σ P[i][j] * Y[i][j]用矩阵运算写就是 x_hat sum(P ⊙ X)y_hat sum(P ⊙ Y)其中 ⊙ 是逐元素乘法。这个计算本质上是求概率分布下坐标的期望值。因为 P 是概率分布X 和 Y 是坐标值所以期望就是加权平均。整个过程是线性的梯度可以无损回传。3.3 为什么用 L2 归一化而不是 softmax这是 DSNT 最核心的设计决策值得单独讲清楚。softmax 的形式是 P[i][j] exp(H[i][j]) / Σ exp(H[a][b])。它的问题在于指数运算会非线性地放大差异。假设热力图上有两个峰值一个响应是 5一个是 3softmax 之后可能变成 0.88 和 0.12差距被拉大了。如果响应是 10 和 8softmax 之后可能变成 0.88 和 0.12 差不多因为指数把绝对差异放大了但相对差异被压缩了。这种非线性变换会改变分布的形状导致期望坐标偏离真实峰值。L2 归一化除以总和是线性的它只是把整个热力图等比例缩放不改变各位置之间的相对关系。所以期望坐标能更忠实地反映热力图的原始分布。这就是 DSNT 在分布保真度上优于 soft-argmax 的根本原因。当然L2 归一化也有它的代价。如果热力图本身比较平坦归一化后分布会很分散期望坐标可能落在几个峰值的中间而不是任何一个峰值上。这个问题在训练初期比较常见因为网络还没学会输出尖锐的热力图。解决办法通常是加一个正则项鼓励热力图集中或者用温度系数调整分布的尖锐程度。3.4 损失函数的设计DSNT 的损失函数通常由两部分组成坐标损失和热力图正则损失。坐标损失就是预测坐标和真实坐标之间的欧氏距离或者 L1 距离L_coord || (x_hat, y_hat) - (x_gt, y_gt) ||热力图正则损失是为了让热力图更接近一个理想的高斯分布避免分布过于分散或者出现多峰。常用的做法是计算预测热力图和真实高斯热力图之间的 KL 散度或者 MSEL_reg KL(P || G) 或 MSE(H, G)总损失是两者的加权和L L_coord λ * L_regλ 是个超参数控制正则强度。我实测下来λ 取 0.1 到 1.0 之间比较合适。太小了热力图会散太大了坐标精度会受影响因为模型会过度关注热力图形状而忽略坐标误差。这个参数需要根据具体任务调后面实操部分会讲怎么调。4. 代码实现与关键细节4.1 核心层的实现DSNT 的核心实现非常简洁用 PyTorch 写大概就几十行。下面是我常用的一个版本加了详细注释import torch import torch.nn as nn import torch.nn.functional as F class DSNT(nn.Module): def __init__(self, height, width, normalizeTrue): super().__init__() self.height height self.width width self.normalize normalize # 构造固定的坐标矩阵 # X 矩阵列索引归一化到 [0, 1] x torch.linspace(0, 1, width).view(1, 1, 1, width) x x.expand(1, 1, height, width) # Y 矩阵行索引归一化到 [0, 1] y torch.linspace(0, 1, height).view(1, 1, height, 1) y y.expand(1, 1, height, width) self.register_buffer(x_mat, x) self.register_buffer(y_mat, y) def forward(self, heatmap): # heatmap: [B, K, H, W] B, K, H, W heatmap.shape # 非负化取平方放大高响应区域 hm heatmap ** 2 # 归一化除以总和得到概率分布 hm hm / (hm.sum(dim(2, 3), keepdimTrue) 1e-8) # 计算坐标期望 x_hat (hm * self.x_mat).sum(dim(2, 3)) # [B, K] y_hat (hm * self.y_mat).sum(dim(2, 3)) # [B, K] # 拼接成坐标 coords torch.stack([x_hat, y_hat], dim-1) # [B, K, 2] return coords, hm这段代码有几个细节值得说。第一register_buffer把坐标矩阵注册成 buffer这样它们会跟着模型一起搬到 GPU但不会被当作可训练参数。第二heatmap ** 2这个非负化操作我试过取绝对值、ReLU、平方几种方案平方的效果最稳定因为它对高响应区域的放大是平滑的不会像 ReLU 那样在零点产生硬截断。第三 1e-8是防止除零虽然理论上概率分布不会全零但训练初期网络输出可能全是负值平方后接近零加个 eps 更保险。4.2 坐标矩阵的构造陷阱坐标矩阵的构造看起来简单但有个容易踩的坑归一化的范围。如果你用torch.linspace(0, 1, width)那坐标范围是 [0, 1]。但如果你希望坐标范围是 [-1, 1]或者直接是像素坐标 [0, width-1]就需要调整。这个选择要和你的 ground truth 坐标范围保持一致否则 loss 会算错。我见过有人热力图坐标矩阵用 [0, 1]但 GT 坐标用的是像素值 [0, 255]结果 loss 一直是几百模型根本学不动。这种低级错误在新手里挺常见的写代码的时候一定要对齐。另外linspace(0, 1, width)在 width1 的时候会返回 [0]这时候除以 (width-1) 会除零。虽然实际任务里热力图尺寸不会小到 1但如果你在写通用库这个边界情况要处理一下。4.3 与网络输出的对接DSNT 层通常接在网络的最后输入是 [B, K, H, W] 的热力图输出是 [B, K, 2] 的坐标。这里有个设计选择热力图是用 sigmoid 激活还是线性输出如果用 sigmoid输出范围是 [0, 1]本身非负DSNT 里可以省掉平方操作直接归一化。但 sigmoid 有个问题它会把输出压缩到 [0, 1]如果网络想表达这个位置响应很强最多只能到 1区分度不够。而且 sigmoid 在饱和区梯度很小训练后期容易梯度消失。如果用线性输出范围不受限表达能力强但需要 DSNT 里做非负化。我倾向于线性输出 平方非负化实测精度更好训练也更稳定。还有一种做法是用 ReLU 输出天然非负但 ReLU 在负值区梯度为零如果某个位置网络输出负值梯度就断了。平方操作在负值区梯度是 2x虽然也是负的但至少不是零信息不会完全丢失。4.4 训练时的损失计算损失计算部分坐标损失我一般用 L1 而不是 L2。L2 对大误差惩罚重容易让模型过度关注难样本训练不稳定。L1 更鲁棒梯度也稳定。当然这个不是绝对的具体任务具体分析。热力图正则损失我用的是预测热力图和真实高斯热力图之间的 MSE。真实高斯热力图的构造方式是以 GT 坐标为中心生成一个二维高斯分布标准差通常取 1 到 2 个像素。标准差太小热力图太尖DSNT 的期望坐标容易受噪声影响标准差太大热力图太平坐标精度上不去。def gaussian_heatmap(coords, height, width, sigma1.5): # coords: [B, K, 2] 归一化坐标 B, K, _ coords.shape y_grid, x_grid torch.meshgrid( torch.arange(height, devicecoords.device), torch.arange(width, devicecoords.device), indexingij ) x_grid x_grid.float() / (width - 1) y_grid y_grid.float() / (height - 1) # 扩展维度便于广播 x_grid x_grid.view(1, 1, height, width) y_grid y_grid.view(1, 1, height, width) x_gt coords[:, :, 0].view(B, K, 1, 1) y_gt coords[:, :, 1].view(B, K, 1, 1) # 高斯分布 gauss torch.exp(-((x_grid - x_gt)**2 (y_grid - y_gt)**2) / (2 * sigma**2)) # 归一化 gauss gauss / (gauss.sum(dim(2, 3), keepdimTrue) 1e-8) return gauss这个高斯热力图的构造sigma 的选择很关键。我一般从 1.5 开始试如果坐标精度不够就减小如果训练不稳定就增大。有些论文会用可学习的 sigma但我觉得没必要固定值够用了还省事。5. 实操流程与调参经验5.1 完整训练流程把 DSNT 集成到你的关键点检测模型里大致流程是这样的网络 backbone 输出特征图经过几个上采样或者反卷积层得到 [B, K, H, W] 的热力图。H 和 W 通常是输入图像的 1/4 或者 1/2K 是关键点数量。热力图经过 DSNT 层得到 [B, K, 2] 的归一化坐标。计算坐标损失预测坐标和 GT 坐标的 L1 距离对所有关键点求平均。计算热力图正则损失预测热力图和 GT 高斯热力图的 MSE。总损失 坐标损失 λ * 正则损失反向传播更新参数。推理阶段直接用 DSNT 输出的坐标不需要再做 argmax 或者后处理。这是 DSNT 的另一个好处训练和推理完全一致没有断层。5.2 学习率和优化器选择DSNT 层本身没有可学习参数所以它不影响优化器的选择。但因为它改变了梯度的传播路径对学习率有一定敏感性。我用 Adam 优化器比较多学习率从 1e-3 开始配合 cosine annealing 或者 step decay。如果发现训练初期 loss 震荡厉害可以把学习率降到 1e-4等稳定了再升回去。SGD 也能用但需要更仔细地调 momentum 和 weight decay不如 Adam 省心。有个经验DSNT 刚接上的时候坐标 loss 可能一开始比较大因为网络还没学会输出合理的热力图。这时候如果学习率太大容易把网络带偏。我通常会用 warmup前几个 epoch 学习率从 1e-5 线性升到 1e-3让网络先适应一下。5.3 λ 参数的调节策略λ 控制坐标损失和热力图正则损失的平衡是 DSNT 里最重要的超参数。λ 太小比如 0.01热力图正则几乎不起作用网络只关注坐标精度热力图可能变得很散或者出现多峰虽然坐标期望可能还对但泛化能力差遇到没见过的样本容易崩。λ 太大比如 10网络过度关注热力图形状坐标精度反而上不去因为热力图像不像和坐标准不准不是完全等价的。我的调节策略是先用 λ1 跑一遍看坐标 loss 和正则 loss 的量级。如果正则 loss 比坐标 loss 大一个数量级以上说明 λ 偏大往下调如果正则 loss 几乎不降说明 λ 偏小往上调。目标是让两个 loss 在同一量级或者正则 loss 略小一点。具体数值上我试过的任务里λ 在 0.1 到 1.0 之间居多。人脸关键点这种热力图比较尖锐的任务λ 可以小一点0.1 到 0.3人体姿态估计这种关节区域比较大的任务λ 可以大一点0.5 到 1.0。5.4 热力图分辨率的取舍热力图的分辨率 H×W 直接影响 DSNT 的精度和计算量。分辨率越高坐标的亚像素精度理论上越好因为期望是在更细的网格上计算的。但计算量也越大而且 DSNT 的坐标期望计算是 O(H×W) 的分辨率翻倍计算量翻四倍。我一般用输入图像的 1/4 分辨率比如输入 256×256热力图 64×64。这个分辨率下DSNT 的亚像素精度已经能到 0.1 像素左右足够大多数任务用了。如果任务对精度要求特别高比如医学图像定位可以用 1/2 分辨率但要注意显存和速度。还有个技巧如果热力图分辨率太低比如 16×16DSNT 的期望坐标可能会被量化误差影响因为网格太粗期望的离散近似误差变大。这种情况下可以在 DSNT 之前加一个上采样层把热力图放大到合适的分辨率再算期望。6. 常见问题与排查技巧6.1 坐标总是偏向图像中心这是 DSNT 最常见的问题尤其是训练初期。原因是热力图分布太散期望坐标被拉向分布的中心而如果热力图没有明显的峰值中心往往就是图像中心。排查思路先可视化热力图看看是不是一片模糊没有明显的峰值。如果是说明网络还没学会输出尖锐的热力图。解决办法有三个一是增大 λ加强热力图正则二是减小高斯热力图的 sigma让目标分布更尖锐三是检查网络结构看看是不是上采样不够特征图分辨率太低导致热力图无法表达细节。我遇到过一次热力图分辨率只有 8×8怎么调 λ 都没用后来把分辨率提到 32×32问题直接消失了。所以分辨率这个基础参数一定要先确认够用。6.2 训练 loss 震荡不收敛DSNT 的 loss 震荡通常和归一化的数值稳定性有关。如果热力图里有很多接近零的值归一化时分母可能很小导致数值不稳定。虽然加了 eps但如果 eps 太小比如 1e-12还是可能出问题。我一般用 1e-8这个量级在 float32 下比较安全。另一个原因是坐标矩阵的范围和 GT 坐标不匹配导致 loss 量级异常。前面提过一定要对齐。检查方法是打印一下预测坐标和 GT 坐标的范围看看是不是在同一量级。还有就是学习率太大。DSNT 的梯度比 argmax 方案要活跃得多因为梯度能回传到每个像素如果学习率太大热力图的每个像素都在剧烈变化loss 自然震荡。降学习率或者加 warmup 通常能解决。6.3 多峰分布导致坐标漂移如果热力图上出现了多个峰值DSNT 的期望坐标会落在这些峰值的加权中心而不是任何一个峰值上。这在关键点被遮挡或者图像模糊时比较常见。解决办法是在损失函数里加一个集中度正则项惩罚分布的熵。熵越小分布越集中多峰的可能性越低。具体做法是在总损失里加一项L_entropy -Σ P[i][j] * log(P[i][j] 1e-8)然后 L L_coord λ1 * L_reg λ2 * L_entropy。λ2 一般取 0.01 到 0.1。不过这个正则项要慎用因为它会鼓励分布变得极尖可能牺牲泛化能力。我一般只在多峰问题比较严重的时候才加而且 λ2 取小一点。6.4 常见问题速查表问题现象可能原因排查方法解决方案坐标偏向中心热力图分布太散可视化热力图增大 λ、减小 sigma、提高分辨率loss 震荡数值不稳定或学习率大检查 eps 和坐标范围增大 eps、对齐坐标范围、降学习率坐标漂移多峰分布检查热力图峰值数量加熵正则、增大 λ精度上不去分辨率太低对比不同分辨率提高热力图分辨率训练不收敛坐标范围不匹配打印坐标范围对齐坐标矩阵和 GT 范围推理速度慢热力图分辨率过高profile 各层耗时降分辨率或加轻量上采样6.5 几个独家避坑技巧第一个技巧训练初期先用较大的 sigma 和较小的 λ让网络先学会大概位置等 loss 降下来再逐渐减小 sigma、增大 λ精调坐标。这个 curriculum learning 的思路比一上来就用最终参数效果好很多。第二个技巧如果任务的关键点数量多比如 68 个人脸关键点DSNT 的坐标期望计算可以并行但热力图正则损失的计算量会比较大。这时候可以只对部分关键点算正则损失或者降低正则损失的采样频率节省计算。第三个技巧DSNT 输出的坐标是归一化的 [0, 1]如果要转成像素坐标记得乘以 (width-1) 和 (height-1)而不是 width 和 height。这个 off-by-one 错误很隐蔽但会影响精度。第四个技巧如果发现某些关键点的精度特别差可以单独可视化这些关键点的热力图看看是不是这些点的分布本身就有问题。有时候是数据标注的问题不是模型的问题。7. 适用场景与扩展思路DSNT 最适合的场景是需要从空间热力图回归精确坐标且对端到端训练有要求的任务。典型的有人体姿态估计17 个或者更多关节点的定位人脸对齐68 点或者 98 点人脸关键点手部关键点21 个手部关节医学图像定位病灶中心、器官关键点目标中心点检测比如 CenterNet 这类方案不太适合的场景如果任务只需要粗略的位置不需要亚像素精度那 argmax 就够了没必要上 DSNT。如果热力图本身不是概率分布的形式比如是分割 mask那 DSNT 也不适用。扩展思路上DSNT 的思想可以推广到三维。把二维的坐标矩阵换成三维的热力图变成体素就能做三维关键点检测。计算量会大很多但原理是一样的。另外DSNT 也可以和注意力机制结合用注意力图代替热力图做可微的注意力定位。我在实际项目里用 DSNT 替换 argmax 之后人脸关键点的平均误差从 3.2 像素降到了 2.1 像素提升还是很明显的。而且训练过程更稳定不需要在推理阶段做额外的后处理整个 pipeline 简洁了不少。如果你也在做类似的任务建议花半天时间把 DSNT 接进去试试大概率不会让你失望。