- 计算机视觉
- 人工智能
- 深度学习
- 图像处理
【免费下载链接】kornia
🐍 Geometric Computer Vision Library for Spatial AI
导读
本文围绕 Kornia 版本迁移记录 changelog.d/+migration-101.fixed.md 展开,深入解析一个与 Apple MPS 后端直接相关的数据增强缺陷:RandomTransplantation与RandomTransplantation3D在未指定excluded_labels时,会在 MPS 设备上"移植失败"——输出完全等于输入。读完本文,你将掌握该 Bug 的根因(PyTorch MPS 后端对空轴all()的未定义求值)、官方修复策略("无物可排除时跳过过滤"),以及如何通过源码与测试用例验证修复行为。
一、Bug 现象与影响范围
1.1 现象描述
根据 changelog.d/+migration-101.fixed.md 的原始记录:
RandomTransplantationandRandomTransplantation3Dtransplanted nothing on MPS when noexcluded_labelswere given.
也就是说,在 Apple MPS(Metal Performance Shaders)后端上,只要用户没有显式传入excluded_labels参数,移植(transplantation)增强就会"静默失效":本应把批次中某张图的某个语义标签对应的像素区域复制粘贴到另一张图上,结果输出与输入完全一致,增强形同虚设。
1.2 影响对象
涉及两个公开 API,均已在 kornia/augmentation/init.py 的__all__中导出("RandomTransplantation"、"RandomTransplantation3D",见 kornia/augmentation/init.py):
RandomTransplantation:2D 版本,定义于 kornia/augmentation/_2d/mix/transplantation.py,继承自MixAugmentationBaseV2;RandomTransplantation3D:3D 版本,定义于 kornia/augmentation/_3d/mix/transplantation.py,多重继承RandomTransplantation与AugmentationBase3D,用于在AugmentationSequential容器内处理(B, C, D, H, W)体积数据。
两者共享同一套核心逻辑,因此 Bug 同时影响 2D 与 3D 场景。该增强算法源自论文《Semantic segmentation of surgical hyperspectral images under geometric domain shifts》(sellner2023semantic),常用于医学影像等需要"跨样本语义对象复制"的场景。
二、RandomTransplantation 的工作原理回顾
要理解这个 Bug,先要理解移植增强的算法流程。根据 kornia/augmentation/_2d/mix/transplantation.py 的文档字符串,其工作方式分为三步:
- 选定受体(acceptor):根据概率参数
p,从批次中挑选若干图像作为"受体"; - 确定供体(donor):每个受体的供体是批次中位于其下方的一张图,通过循环取模规则
i - 1 mod B得到(即donor_indices = (acceptor_indices - 1) % batch_size,见 params_from_input); - 随机移植:从供体的语义分割掩码中随机选取一个标签(label),把供体上该标签对应的图像特征与分割掩码区域整体复制到受体对应位置。
移植由批次中第一个"mask"输入驱动,可同时作用于图像(DataKey.INPUT,形状(B, C, *spatial))与掩码(DataKey.MASK,形状(B, *spatial))。数据键的默认顺序为[DataKey.INPUT, DataKey.MASK],见init。
2.1 excluded_labels 参数的作用
excluded_labels是移植增强的"黑名单":
sequence of labels which should not be transplanted from a donor. This can be useful if only parts of the image are annotated and the non-annotated regions (with a specific label index) should be excluded from the augmentation.
它常用于"图像只有部分区域被标注"的场景——例如标签 0 代表未标注区域,此时应把标签 0 排除在可移植标签之外,避免把无标注区域复制到其他样本上。构造函数中该参数默认值为None,会被归一化为空列表再转为空张量(见init):
if excluded_labels is None: excluded_labels = [] if not isinstance(excluded_labels, torch.Tensor): excluded_labels = torch.tensor(excluded_labels) self.excluded_labels: torch.Tensor = excluded_labels2.2 标签筛选的底层实现
在 params_from_input 中,供体掩码的唯一标签会与excluded_labels做差集过滤:
if "selected_labels" not in params and "selection" not in params: if self.excluded_labels.device != mask.device: self.excluded_labels = self.excluded_labels.to(mask.device) donor_labels: list[torch.Tensor] = [] eligible: list[int] = [] for d in range(len(params["donor_indices"])): current_mask = mask[params["donor_indices"][d]] labels = current_mask.unique() if self.excluded_labels.numel() > 0: labels = labels[(labels.view(1, -1) != self.excluded_labels.view(-1, 1)).all(dim=0)] ...关键逻辑在于labels.view(1, -1) != self.excluded_labels.view(-1, 1)这行:它构造一个"标签 × 排除项"的布尔矩阵,随后沿排除项轴(dim=0)做.all()归约,从而筛选出"不等于任何排除标签"的候选标签。
三、Bug 根因:MPS 后端对空轴all()的未定义求值
3.1 为什么默认情况下会触发
由于excluded_labels默认为空(numel() == 0),当用户不传该参数时,布尔矩阵的排除项维度大小为 0。此时沿空轴做all()归约,在 CPU 与 CUDA 后端上会返回True(数学上"空集的全称命题为真",即全标签都保留),但在PyTorch 的 MPS 后端(PyTorch 2.9)上,对空轴的all()求值结果是未定义的,通常返回False。
返回值一旦是False,就意味着每一个供体标签都会被过滤掉:
labels被清空为长度为 0 的张量;len(labels) > 0判断失败,该供体被判定为"无合格标签";- 其对应的受体从
acceptor_indices中被剔除(params_from_input),batch_prob对应项被清零; - 结果就是所有受体都不再是受体,输出等于输入。
这正是迁移记录中描述的"transplanted nothing"的完整机制。需要注意,此处all()的语义是沿"排除项"这一空轴归约,而非沿标签轴,因此修复不能简单地替换all()的调用方式。
3.2 为什么这是一个隐蔽的跨后端一致性问题
同类归约在 CPU/CUDA 与 MPS 上的行为不一致,属于典型的**后端语义差异(backend semantic divergence)**问题。同样的代码在 Linux/CUDA 环境测试通过,一旦迁移到 Apple Silicon 的 MPS 环境,增强器会静默失效,且不抛任何异常——这对训练管线极具迷惑性,因为"增强无效果"不会导致报错,只会降低数据多样性、影响模型泛化,且极难定位。
四、官方修复方案:无物可排除时跳过过滤
4.1 修复后的源码
修复的核心思路非常直接:当excluded_labels为空时,完全跳过差集过滤步骤,不做空轴all()归约。修复后的代码在 kornia/augmentation/_2d/mix/transplantation.py 中体现为:
# Remove any label which is part of the excluded labels. Skip the reduction when there is nothing # to exclude: on MPS (PyTorch 2.9) ``all`` over an empty axis yields an undefined result, usually # False, where CPU and CUDA return True, and the filter would then discard every label. if self.excluded_labels.numel() > 0: labels = labels[(labels.view(1, -1) != self.excluded_labels.view(-1, 1)).all(dim=0)]修复后,空排除表场景下labels直接保留current_mask.unique()的全部结果,标签抽取(labels[torch.randperm(len(labels))[0]])得以正常进行,移植逻辑恢复可用。
4.2 修复的边界影响
需要特别指出的是,该修复没有改变以下既有行为(这些行为仍由 kornia/augmentation/_2d/mix/transplantation.py 中的逻辑保证):
- 当供体掩码本身没有任何标签(如全空掩码)时,
len(labels) > 0依旧为假,该供体仍会被判为"无物可给",其受体照常被剔除——这是算法设计的预期行为,与排除表无关; - 当
excluded_labels非空、但恰好排除了供体的全部标签时,过滤逻辑正常运行(此时归约轴非空,MPS 行为确定),受体同样会被剔除; - 标签抽取仍基于全局 CPU 随机生成器的
torch.randperm(测试test_convention_label_draw_uses_the_cpu_generator_whatever_the_device验证了这一点),无论掩码位于何种设备,随机状态不受影响。
五、测试用例如何验证修复
修复并非孤立的代码改动,而是伴随一套可执行测试,用于固化约定并防止回归。围绕移植增强的约定测试集中在 tests/augmentation/test_conventions_transplantation.py 与 tests/augmentation/test_conventions_3d.py 中,相关变更记录见 changelog.d/4695.added.md("Document the conventions ofRandomTransplantationandRandomTransplantation3D... with executable tests")。
与本文 Bug 直接相关的测试覆盖点包括:
- 默认空排除表的行为:例如
test_convention_transplant_follows_the_device_and_dtype_of_its_inputs(tests/augmentation/test_conventions_transplantation.py)在构造RandomTransplantation(p=1.0, excluded_labels=[0])时明确注释 "the exclusion list starts on the CPU",验证排除列表会随输入迁移到目标设备,且移植结果等于整行滚动image.roll(1, dims=0); - 空掩码的边界情况:测试注释明确 "An empty mask has no label at all, so nothing is eligible: empty in, empty out, no raise"(见 tests/augmentation/test_conventions_transplantation.py),说明空供体掩码的"无合格标签"路径与空排除表的路径是两条独立逻辑,前者仍按设计剔除受体;
- 供体无需是受体:
test_convention_donor_need_not_be_an_acceptor验证供体取完整批次中的前驱(i - 1) mod B,即使该前驱本身未被选中为受体也照常充当供体(tests/augmentation/test_conventions_transplantation.py); - 3D 类的容器行为:tests/augmentation/test_conventions_3d.py 枚举 3D 子类,确认
RandomTransplantation3D是容器中唯一携带 mixinverse(抛RuntimeError)的 3D 增强,并验证p_batch=0.0时完全跳过。
这些测试同时锁定了_params中batch_prob、acceptor_indices、donor_indices、selected_labels、selection等字段的契约,确保后续任何后端相关的改动不会破坏参数重放(replay)语义。
六、对使用者的实操建议
结合本次修复,在实际使用RandomTransplantation/RandomTransplantation3D时有几点值得注意:
- 升级即修复:该修复随 migration-101 合入当前仓库主干。若你在 Apple Silicon + MPS 环境下使用 Kornia,且观察到移植增强"无效"(输出等于输入),升级到包含 #4160 修复的版本即可解决;
- 默认参数是安全的:修复后,不传
excluded_labels(即默认空表)在 MPS、CPU、CUDA 上行为一致,无需为了规避 Bug 而刻意传入空列表或占位排除项; - 显式排除表的语义不变:如果你用
excluded_labels排除未标注区域(如标签 0),请确保排除列表确实与掩码中的标签值对应;一旦供体掩码中剩余标签为空,对应受体仍会被剔除(这是设计行为,见 params_from_input),批次其余部分照常移植; - 注意容器中的排名约束:在
AugmentationSequential中使用时,2D 图像用RandomTransplantation、5D 体积用RandomTransplantation3D;且移植类在容器中只能作为第一步——一旦其他增强运行过,容器会把掩码递进为(B, 1, H, W),移植类会因排名规则拒绝处理(详见 kornia/augmentation/_2d/mix/transplantation.py 的约定说明与 issue #4707); - 跨设备迁移注意:
excluded_labels张量初始创建在 CPU(构造函数中torch.tensor(excluded_labels)),会在params_from_input中自动迁移到掩码所在设备(kornia/augmentation/_2d/mix/transplantation.py),无需手动干预。
七、总结
本次 #4160 修复是典型的"后端语义差异导致的静默失效"案例:PyTorch MPS 后端(2.9)对空轴all()返回未定义值(通常False),而 CPU/CUDA 返回True,导致RandomTransplantation/RandomTransplantation3D在默认无excluded_labels时过滤掉全部供体标签,输出等于输入。修复方案简单而精准——"无物可排除时跳过过滤归约",配合 tests/augmentation/test_conventions_transplantation.py 与 tests/augmentation/test_conventions_3d.py 中的可执行约定测试,保证了跨后端行为的一致性。对于在 Apple Silicon 上做语义分割数据增强的开发者,这一修复直接关系到移植增强是否真正生效,值得纳入升级清单。
- 计算机视觉
- 人工智能
- 深度学习
- 图像处理
【免费下载链接】kornia
🐍 Geometric Computer Vision Library for Spatial AI
相关推荐
Kornia 修复 RandomTransplantation 在 MPS 上的移植失效:空 excluded_labels 触发的 all() 归约陷阱与规避方案
Kornia 修复 RandomTransplantation 在 MPS 上的移植失效:空 excluded_labels 触发的 all 归约陷阱与规避方案
计算机视觉深度学习人工智能图像处理Kornia 修复 Apple MPS 后端 SVD 8192 元素上限:`_torch_svd_cast` 的 CPU 回退机制解析
Kornia 修复 Apple MPS 后端 SVD 8192 元素上限: _torch_svd_cast 的 CPU 回退机制解析 导读 本文围绕 Korni
计算机视觉人工智能深度学习图像处理Kornia 修复 MPS 等加速器后端的空张量转换:YUV420/YUV422 空输入语义与 reshape 歧义问题深度解析
Kornia 修复 MPS 等加速器后端的空张量转换:YUV420/YUV422 空输入语义与 reshape 歧义问题深度解析 导读 :本文围绕 Kornia
计算机视觉深度学习人工智能图像处理
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考