Kornia 修复 RandomThinPlateSpline(scale=0.0) 退化分布报错：零控制点位移的恒等变换实现解析

本篇文章围绕 Kornia 仓库中 changelog.d/+migration-013.fixed.md 记录的修复展开：RandomThinPlateSpline(scale=0.0) 此前会在构造退化均匀分布 Uniform(0, 0) 时抛出错误，修复后改为直接生成零控制点位移（即恒等变换）。读完本文，你将理解该增广的薄板样条（Thin Plate Spline, TPS）原理、零尺度分支的源码实现、测试验证方式，以及升级后对既有代码的迁移影响。

一、修复背景：RandomThinPlateSpline 是什么

RandomThinPlateSpline 是 Kornia 2D 几何增广模块（kornia/augmentation/_2d/geometric/thin_plate_spline.py）中基于薄板样条插值实现的随机形变增广，被封装为 AugmentationBase2D 的子类，可直接用于数据增强流水线。它内部依赖 kornia/geometry/transform/thin_plate_spline.py 中的 get_tps_transform 与 warp_image_tps 完成坐标变换与图像重采样，并已从 kornia.augmentation、kornia.augmentation._2d.geometric 以及容器模块 kornia/augmentation/container/augment.py 多处导出，属于公开 API 的一部分。

该增广的约定（见类文档字符串）为：使用 5 个归一化坐标控制点 (-1, -1)、(-1, 1)、(1, -1)、(1, 1)、(0, 0)，对每个目标坐标施加 [-scale, scale] 区间的均匀噪声；默认采用双线性采样、零填充（padding_mode='zeros'）与 align_corners=False，且不暴露插值模式参数。其构造参数如下：

参数	默认值	说明
scale	0.2	作用于目标点的非负噪声尺度，为 0 时控制点保持不变
align_corners	False	传给 grid_sample 的对齐标志
padding_mode	'zeros'	传给 grid_sample 的填充方式：'zeros'、'border' 或 'reflection'
same_on_batch	False	是否对整批样本应用相同变换
p	0.5	应用该变换的概率
keepdim	False	输出保持输入形状（True）或广播为批形式（False）

二、问题根源：Uniform(-scale, scale) 的退化分布

在修复之前，构造逻辑会在 scale 取任意值时都创建一个 torch.distributions.Uniform(-scale, scale) 实例用于采样控制点噪声：

self.dist = torch.distributions.Uniform(-scale, scale)

当用户显式传入 scale=0.0 时，Uniform(0, 0) 是一个退化分布（上下界相等）。按照 torch.distributions 的实现，这种退化分布会触发校验错误，从而在构造 RandomThinPlateSpline(scale=0.0) 时直接抛出异常。

这在功能上并不合理：scale=0 在语义上应当表示"控制点不做任何位移"，即薄板样条退化为恒等变换。用户可能出于以下动机显式传入 0.0：

• 在超参数网格搜索或配置文件中将"形变强度"置零，作为无增强的对照实验；
• 调试时临时关闭形变，只观察其他增广的组合效果；
• 编写自动化测试，验证零强度下增广保持输入不变。

旧行为将这些合理用法挡在门外，因此被列为 bug 并在 changelog.d/+migration-013.fixed.md 中以 fixed 类别记录（关联 issue #4463）。

三、修复方案：零尺度直接生成零位移

修复后的关键改动是：当 scale == 0 时不再构造任何分布对象，而是在参数生成阶段直接输出与源控制点完全一致的零位移目标点。对应源码见 kornia/augmentation/_2d/geometric/thin_plate_spline.py：

self.dist = None if scale == 0 else torch.distributions.Uniform(-scale, scale)

在 generate_parameters 中，dist is None 的分支直接用全零张量代替采样：

# 5 TPS control points in normalized coordinates
src = _constant_tensor(
    [[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0], [0.0, 0.0]]],
    device=device,
    dtype=dtype,
).expand(B, 5, 2)

if self.dist is None:
    noise = torch.zeros_like(src)
elif self.same_on_batch:
    noise = self.dist.rsample((1, 5, 2)).to(device=device, dtype=dtype)
    noise = noise.expand(B, 5, 2)
else:
    noise = self.dist.rsample((B, 5, 2)).to(device=device, dtype=dtype)

dst = src + noise

可以看到：

1. 源控制点 src 的归一化坐标为 (-1,-1), (-1,1), (1,-1), (1,1), (0,0)，形状为 (B, 5, 2)；
2. scale=0 时噪声为全零张量，因此 dst == src；
3. 非零尺度且 same_on_batch=True 时只采样一份噪声再 expand 到整批，保证批内一致；same_on_batch=False 时每张图独立采样；
4. 采样使用 rsample（重参数化采样），保持梯度可回传，这对后续反向传播与 gradcheck 很重要。

在 apply_transform 中，变换通过"逆参数"方式计算——即调用 get_tps_transform(dst, src)（第二个位置参数作为核中心）得到核权重与仿射权重，再交给 warp_image_tps 完成从输出空间到输入空间的采样：

kernel, affine = get_tps_transform(dst, src)
return warp_image_tps(input, src, kernel, affine, flags["align_corners"], flags["padding_mode"].name.lower())

四、恒等性质的数学与数值说明

当 dst == src 时，TPS 的核权重在数学上为零、仿射部分为单位映射，整体应等价于恒等变换。但需要注意，get_tps_transform 内部通过解线性方程组（kornia/geometry/transform/thin_plate_spline.py 中 _pair_square_euclidean 构造核矩阵后使用线性求解器）得到权重，因此恒等映射只在线性求解器的舍入误差范围内实现，并非逐位精确相等，残余大小取决于 dtype 与后端。

当前类文档明确承诺：在 float16、float32、float64 与 bfloat16 下，零尺度的图像 warp 在任一 align_corners 设置下都是"数值精度内的恒等变换"。从源码结构看，这也是本修复选择"零位移"而非"仍然采样后取近似"的直接原因——scale=0 分支完全绕开了随机采样，使得结果确定且可复现，同时避免了退化分布构造异常。

五、测试验证：从参数到梯度逐层锁定

仓库测试为该修复提供了完整的验证证据，读者可自行复跑确认：

1. 零尺度参数生成（tests/augmentation/test_augmentation.py 中 TestRandomThinPlateSpline.test_zero_scale_parameters）：对 batch_size ∈ {0, 1, 4} 与 same_on_batch ∈ {False, True} 的全部组合，断言 scale=0.0 时生成的 dst 形状为 (batch_size, 5, 2)，设备与 dtype 正确，且与 src 在 atol=0, rtol=0 下完全相等——这是对"零控制点位移"最直接的回归测试。

2. 零尺度恒等性（同文件 test_zero_scale_identity）：对两种 align_corners 与两种 same_on_batch 组合，构造一个 6×8 的梯度图像后经 RandomThinPlateSpline(scale=0.0) 前向传播，断言输出与输入逐元素相等、使用固定 params 重放的结果一致，且 output.sum().backward() 后输入梯度为全 1。这同时验证了零尺度分支的可微性与梯度流。

3. 约定一致性（tests/augmentation/test_conventions_geometric_dense.py）：test_convention_tps_control_points_and_noise_sample_on_cpu 验证 scale=0 时 src == dst、scale=0.2 时 |dst - src| <= 0.201（保证噪声不越界），以及默认 padding_mode == SamplePadding.ZEROS、align_corners is False；test_convention_thin_plate_spline_identity_grid_3928 则对 align_corners ∈ {False, True} 直接断言输出等于输入，覆盖了历史 issue #3928 的采样网格对齐问题。

4. 已知边界：test_zero_scale_identity 中 float16 被显式跳过，注释指出 get_tps_transform 在 float16 下数值不稳定（可能产生 NaN）。这意味着虽然文档承诺 bfloat16 等精度可用，但 float16 下的 TPS 恒等保证并不成立，属于已知限制而非本修复引入的回归。

六、使用指南与迁移影响

修复后，RandomThinPlateSpline(scale=0.0) 成为一个合法的"空操作"增广，可以直接运行：

import torch
import kornia.augmentation as K

img = torch.randn(2, 3, 64, 64)

# 修复前：构造即抛错；修复后：等价于恒等变换
aug = K.RandomThinPlateSpline(scale=0.0, p=1.0)
out = aug(img)
assert out.shape == img.shape  # 输出形状不变

# 多次前向（默认不固定参数）输出一致，因为零尺度下不存在随机采样
torch.testing.assert_close(out, img)

结合其他特性使用时的注意事项：

• 接口限制：RandomThinPlateSpline 不提供 transform_matrix 与 inverse 接口（见类文档与测试 test_convention_dense_shape_and_matrix_interface），因此它无法纳入基于矩阵的增广推理流程，也不支持逆变换；在容器中使用空间标签时另有额外限制。
• 参数重放：与所有 Kornia 增广一致，可通过 aug(input, params=aug._params) 复现同一次变换（默认 p=0.5 时需先捕获 _params）。零尺度分支由于无随机性，天然满足可复现性。
• 迁移影响：此次修复属于 bug fix（changelog 类别为 fixed），不改变 RandomThinPlateSpline 的公开签名与默认行为（默认 scale=0.2 的随机路径不受影响）。唯一行为变化是 scale=0.0 从"构造报错"变为"恒等变换"，因此依赖异常来拦截零强度配置的代码需要调整为显式判断；而在配置文件中合法使用 scale=0.0 的代码则可以直接升级，无需任何改动。
• 数值前提：恒等保证建立在"数值精度内"这一前提下，且不适用于 float16；如需严格恒等，建议在 float32 及更高精度下使用。