Kornia 小尺寸分形生成修复：diamond_square 在 1px/2px 图片上的边界问题与 3px 网格回退方案

Kornia 小尺寸分形生成修复：diamond_square 在 1px/2px 图片上的边界问题与 3px 网格回退方案

【免费下载链接】kornia🐍 Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/GitHub_Trending/ko/kornia

导读

Kornia 的kornia.contrib.diamond_square基于经典 Diamond-Square 算法生成等离子分形图像，并支撑RandomPlasmaBrightness、RandomPlasmaContrast、RandomPlasmaShadow三款 2D 强度增强。当输入图片的空间边长低于 3 像素时，该算法此前会抛出ValueError: math domain error（1px）或TypeError（2px），导致整条数据增强管线在这些极小图上崩溃。本文以仓库变更记录 changelog.d/+migration-004.fixed.md 为主线，结合 diamond_square 源码 与 回归测试，完整还原故障根因、修复策略（3px 最小网格生成 + 切片裁剪）以及 bit-identical 兼容性保证，并给出受影响 API 的参数速查与实战使用建议。

一、问题背景：分形增强在小尺寸输入上集体崩溃

Diamond-Square 是一种递归细分算法，其最小可构造网格单元是 3×3——即两个角点加一个中点。Kornia 将其封装为kornia.contrib.diamond_square（位于 kornia/contrib/diamond_square.py），并在 kornia/augmentation/_2d/intensity/plasma.py 中基于它实现了三个增强类：

三者均在generate_parameters阶段调用_generate_plasma生成整图分形图，再在apply_transform中按输入尺寸裁剪（_plasma_for_input）。这意味着只要输入图片任一边长小于 3px，分形图生成就会失败，进而拖垮三个增强类——这正是本次修复要解决的核心问题。

二、故障现象与根因分析

变更记录原文精确描述了两种失败模式：

1 px side raisedValueError: math domain errorand a 2 px side aTypeErrorfrom a float in thetorch.randsize。

2.1 1px 输入：ValueError: math domain error

从源码看，diamond_square在计算最小 2 的幂网格时执行：

p2_width: int = 2 ** math.ceil(math.log2(max(width, 3) - 1)) + 1

在修复前，这里的max(width, 3)尚不存在（直接使用原始width）。当width == 1时，width - 1 == 0，math.log2(0)在数学上无定义，math模块直接抛出ValueError: math domain error。同理，高度为 1 时也会在p2_height处触发。

2.2 2px 输入：TypeError（float 传入torch.randsize）

当width == 2时，width - 1 == 1，math.log2(1) == 0，这一层不再报错；但在_diamond_square_seed（kornia/contrib/diamond_square.py#L33-L87）内部，随机数采样执行：

res: torch.Tensor = random_fn([replicates, 1, width, height], device=device, dtype=dtype)

在旧的调用路径下，部分尺寸参数由幂运算与除法推导而来（形如(height + 1) // 2的浮点中间值），一旦某个 size 分量是float而非int，torch.rand就会抛出TypeError，即变更记录中所说的 "aTypeErrorfrom a float in thetorch.randsize"。

2.3 影响范围

上述两个异常在 1px/2px 边长的输入上必然触发，因此：

• RandomPlasmaBrightness、RandomPlasmaContrast、RandomPlasmaShadow直接对1×1、1×N、2×2等极小图做增强时全部失败；
• 修复前没有任何合法的降级路径，属于「小尺寸输入即崩溃」的硬边界问题。

三、修复方案：3px 最小网格生成 + 切片裁剪

变更记录给出的修复策略是一句话的核心：小边长先在 3px 网格上绘制（drawn on the 3 px grid），再切片（sliced）得到目标尺寸；3px 及以上的输出与修复前逐位一致（bit-identical）。

这一策略在当前源码中有三处直接体现（见 kornia/contrib/diamond_square.py#L209-L229）：

① 计算 2 的幂边长时，以max(size, 3)作为下限：

p2_width: int = 2 ** math.ceil(math.log2(max(width, 3) - 1)) + 1 p2_height: int = 2 ** math.ceil(math.log2(max(height, 3) - 1)) + 1

这样即使请求 1px 或 2px，内部计算也会落在合法域内（log2的参数至少为 2），从根上消除了math domain error。

② 递归深度与种子网格以 3px 网格为最小单元：

recursion_depth: int = int(min(math.log2(p2_width - 1) - 1, math.log2(p2_height - 1) - 1)) seed_width: int = (p2_width - 1) // 2**recursion_depth + 1 seed_height: int = (p2_height - 1) // 2**recursion_depth + 1

对应的_diamond_square_seed内部也有显式断言KORNIA_CHECK(width == 3 or height == 3, "Height or Width must be equal to 3.")，保证种子网格永远是 3×N 或 N×3 的合法形状，所有参与torch.rand的尺寸都是整数，彻底规避TypeError。

③ 最终按目标尺寸切片：

# slice to match with the output size img = img[..., :width, :height] img = img.view(output_size)

分形图先生成完整 3px 网格及其递归放大结果，再通过img[..., :width, :height]裁出 1px/2px 的目标区域——小尺寸的输出等价于「大图左上角的一个角点/一条边」，与 Diamond-Square 的角点采样语义一致。

3.1 bit-identical 保证的含义

变更记录强调：sides of 3 px and more are bit-identical。这意味着本次修复只影响 1px/2px 边界场景，对 3px 及以上的所有输入，算法路径、随机采样顺序、数值计算完全不变，已有用户不会观察到任何输出变化。这是一次严格向后兼容的缺陷修复，而非行为变更。

四、回归测试：边界尺寸矩阵验证

修复附带了专门的回归测试，见 tests/contrib/test_diamond_square.py：

@pytest.mark.parametrize("spatial_size", [(1, 1), (2, 2), (2, 3), (3, 2), (1, 5), (2, 17)]) def test_spatial_sizes_below_three(self, spatial_size, device, dtype): # Sides under 3 px used to fail with `math domain error` or a float in the `rand` size # (kornia#4472). torch.manual_seed(0) output_size = (2, 3, *spatial_size) out = kornia.contrib.diamond_square(output_size, device=device, dtype=dtype) assert out.shape == output_size assert torch.isfinite(out).all()

该测试覆盖了以下维度组合：

• 双轴都小于 3：(1, 1)、(2, 2)；
• 单轴小于 3、另一轴为 3：(2, 3)、(3, 2)——分别对应_diamond_square_seed中height == 3时触发的转置分支；
• 单轴小于 3、另一轴任意：(1, 5)、(2, 17)，验证切片逻辑在非对称尺寸下同样成立。

断言不仅检查输出形状与请求一致，还要求torch.isfinite(out).all()，即所有输出不含 NaN/Inf——这直接回击了此前math domain error可能留下的非有限值隐患。

与此同时，tests/contrib/test_diamond_square.py#L61-L64 的test_plasma_augmentation_on_a_two_pixel_image还专门验证了增强层的端到端可用性：

def test_plasma_augmentation_on_a_two_pixel_image(self, device, dtype): img = torch.rand(1, 1, 2, 2, device=device, dtype=dtype) out = kornia.augmentation.RandomPlasmaBrightness(p=1.0)(img) assert out.shape == img.shape

而在 tests/augmentation/test_augmentation.py#L5254-L5374 的TestRandomPlasma中，还针对三个增强类覆盖了same_on_batch地图复用（plasma == plasma[:1].expand_as(plasma)）、逐样本独立地图、params=参数重放位级一致性等行为，证明修复没有破坏增强层既有的参数化与可复现特性。

五、实战使用：参数说明与边界注意事项

修复后，diamond_square可以安全接受任意>= 1的空间尺寸，其完整签名与参数如下（kornia/contrib/diamond_square.py#L159-L167）：

diamond_square( output_size: Tuple[int, int, int, int], # (B, C, H, W)，必须为 4 元组 roughness: Union[float, torch.Tensor] = 0.5, # 每次递归的缩放系数 random_scale: Union[float, torch.Tensor] = 1.0, # 递归初始随机幅度 random_fn: Callable[..., torch.Tensor] = torch.rand, # 随机采样函数 normalize_range: Optional[Tuple[float, float]] = None, # 非 None 时做 min-max 归一化到该区间 device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, ) -> torch.Tensor # 输出形状 (B, C, H, W)

直接使用的示例（含归一化，来自 tests/contrib/test_diamond_square.py#L37-L49）：

import torch import kornia torch.manual_seed(0) out = kornia.contrib.diamond_square( output_size=(1, 1, 3, 4), roughness=0.5, random_scale=1.0, normalize_range=(0.0, 1.0), ) assert out.min() == 0.0 and out.max() == 1.0 # 归一化后严格落在区间端点

对应三个增强类的使用方式（参考 plasma.py 中的 docstring 示例）：

>>> rng = torch.manual_seed(0) >>> img = torch.ones(1, 1, 3, 4) >>> RandomPlasmaBrightness(roughness=(0.1, 0.7), p=1.)(img) tensor([[[[0.6415, 1.0000, 0.3142, 0.6836], [1.0000, 0.5593, 0.5556, 0.4566], [0.5809, 1.0000, 0.7005, 1.0000]]]])

使用时的注意事项（均来自类 docstring 的 Convention/Warning 块）：

• 输出始终clamp到[0, 1]：即使输入越界，三个增强类也会将结果限制在该区间；
• 极小尺寸输入现在安全：三个类都声明 "a one-pixel spatial axis is accepted, down to a1 x 1image"，即 1×1 输入也可用；
• 已知限制（issue #4430）：当输入全为负值且采样到的分形图无法将其抬升时，输出可能变成全零图；RandomPlasmaShadow在默认非正shade_intensity下同样存在该现象，可改用正的shade_intensity缓解；
• same_on_batch=True时的地图复用：RandomPlasmaShadow共享的是单通道掩膜（对样本所有通道生效），另外两个类共享(B, C, H, W)全通道地图；跨批次时通过expand广播，保证同批输出逐位一致（对应测试 tests/augmentation/test_augmentation.py#L5282-L5293）；
• 随机种子可复现：diamond_square与增强类的输出均依赖torch.manual_seed，配合params=重放可得到逐位相同的增强结果。

六、小结

本次修复（变更记录 changelog.d/+migration-004.fixed.md，关联 issue #4472、#4488）以极小的改动面解决了三个 Plasma 增强类在 1px/2px 输入上的系统性崩溃：根因是math.log2在 0 处的数学域错误与浮点尺寸传入torch.rand的类型错误；修复采用「以 3px 为最小生成网格、按需切片」的策略，同时通过max(width, 3)与img[..., :width, :height]保证 3px 以上输出逐位不变。配套的 diamond_square 测试 用参数化尺寸矩阵验证了所有边界组合，是「小缺陷 + 严格回归 + 兼容性保证」的一个典型范例，也为在目标检测、超分等场景中处理极小 ROI 图像的增强需求扫清了障碍。

【免费下载链接】kornia🐍 Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/GitHub_Trending/ko/kornia