Kornia mix 增强的 bfloat16 支持与半精度 dtype 保持MixUp / CutMix 实现解析【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/gh_mirrors/ko/kornia本篇文章聚焦 Kornia 增强模块kornia.augmentation的一项近期修复为全部 mix 类增强MixAugmentationBaseV2 家族补齐 bfloat16 支持并让 RandomMixUpV2 与 RandomCutMixV2 在输入为 float16 / bfloat16 时保持输入精度、同时把标签安全提升为 float32。读完本文你将掌握 mix 增强家族的 dtype 处理约定、源码中的核心实现路径以及如何在训练管线中安全使用半精度 MixUp/CutMix。变更背景mix 增强家族与精度丢失问题Mix 类增强mix augmentation与几何增强不同——它不做仿射变换、没有变换矩阵、也不支持inverse()参见 mix 基类约定而是通过样本混合实现正则化当前仓库中属于该家族的有RandomMixUpV2整图线性混合实现RandomCutMixV2按 bbox 裁剪区域粘贴实现RandomJigsaw、RandomMosaic、RandomPatchMix、RandomTransplantation等其余 mix 算子在本次修复之前MixUp/CutMix 处理标签时会把标签强制转换为图像自身的 dtype。这在 float32/float64 输入下没有问题但一旦输入是半精度张量就会踩坑bfloat16尾数只有 8 位float16尾数也只有 11 位而分类任务中类别 id 往往大于 256标签被 cast 到半精度后大整数被截断例如类别 id257在 bfloat16 下会变成256、999会变成1000——这正是测试注释中明确指出的回归现象见 test_conventions_mix.py 中#4657issue 的记录。同时mix 类增强此前不接受 bfloat16 输入与 Kornia 对半精度AMP 训练的支持趋势脱节。changelog.d/4471.fixed.md记录的本项变更即包含两个动作为 mix 增强增加 bfloat16 支持以及在 MixUp/CutMix 中保持输入半精度 dtype、标签提升为 float32 以保证精确。变更一mix 增强基类统一接纳 bfloat16所有 mix 算子共享抽象基类MixAugmentationBaseV2kornia/augmentation/_2d/mix/base.py。dtype 白名单在transform_tensor中统一收口# kornia/augmentation/_2d/mix/base.py#L111-L120 def transform_tensor(self, input, *, shapeNone, match_channelTrue): _validate_input_dtype(input, accepted_dtypes[torch.bfloat16, torch.float16, torch.float32, torch.float64]) if shape is None: return _transform_input(input) return _transform_input_by_shape(input, reference_shapeshape, match_channelmatch_channel)要点白名单同时包含torch.bfloat16与torch.float16意味着RandomMixUpV2 / RandomCutMixV2 / RandomJigsaw / RandomMosaic 等全部 mix 算子都接受 bfloat16 输入该校验同样在**参数回放replay**路径中生效forward在params非空时仍会对输入执行_validate_input_dtypebase.py#L282-L291所以回放一个来自 float32 采样、却在 bfloat16 张量上重放的参数字典会直接抛出TypeError整数输入uint8、int64、bool依旧被拒绝测试test_convention_mix_rejects_integer_images与test_convention_mix_replay_rejects_an_unsupported_dtype_with_type_error覆盖了这一点test_conventions_mix.py、#L589-L597。基类会为每次前向在self._params中记录一个dtype字段其数值来自DType枚举的映射kornia/constants.pyclass DType(Enum, metaclass_KORNIA_EnumMeta): INT64 0 FLOAT16 1 FLOAT32 2 FLOAT64 3 BFLOAT16 4 # 与 torch.bfloat16 一一对应变更二MixUp 图像保持输入 dtype、标签提升为 float32RandomMixUpV2的图像混合在apply_transform中完成核心一行是# kornia/augmentation/_2d/mix/mixup.py#L127-L130 input_permute input.index_select(dim0, indexparams[mixup_pairs].to(input.device)) lam params[mixup_lambdas].view(-1, 1, 1, 1).expand_as(input).to(input.device, dtypeinput.dtype) return input * (1 - lam) input_permute * lam注意lam被显式to(dtypeinput.dtype)混合系数与混合结果都保持在输入的 bfloat16/float16 精度上图像分支不会再被悄悄提升为 float32 后又转回半精度。标签分支则不同。apply_transform_class/apply_non_transform_class使用统一的calc_dtype规则# kornia/augmentation/_2d/mix/mixup.py#L151-L152 image_dtype DType.to_torch(int(params[dtype].item())) calc_dtype image_dtype if image_dtype in (torch.float32, torch.float64) else torch.float32即float32/float64 输入 → 标签沿用图像 dtypefloat16/bfloat16 输入 → 标签一律提升为 float32。最终输出三列(B, 3)torch.stack([input, labels_permute, params[mixup_lambdas]], -1)三列分别为原始标签、配对标签、混合系数 lambda。提升为 float32 的原因在类 docstring 中有明确约定保证整数类别 id 最大到2 ** 24依然精确mixup.py#L78-L83。lambda 的采样逻辑由MixupGenerator负责——默认在[0.0, 1.0]上均匀采样并受_joint_range_check约束random_generator/_2d/mixup.py。变更三CutMix 的同一套 dtype 约定RandomCutMixV2遵循完全相同的约定cutmix.py#L73-L80图像分支apply_transform用 bbox 掩码做区域替换输入为半精度时全程保持在半精度cutmix.py#L186-L199标签分支输出形状为(num_mix, B, 3)每列同样是原始标签 / 配对标签 / lambda且半精度输入下标签提升为 float32cutmix.py#L137-L163lambda 由裁剪区域面积占比推导lam_val w * h / (width * height)use_correct_lambdaTrue时返回1 - lam_val与 CutMix 原论文一致默认False时返回lam_val并发出DeprecationWarningcutmix.py#L127-L135。测试验证精确性与 dtype 保持的回归防线该变更在 tests/augmentation/test_conventions_mix.py 中有系统性验证其中与本次修复最直接的是# tests/augmentation/test_conventions_mix.py#L538-L558 pytest.mark.parametrize(image_dtype, [torch.float16, torch.bfloat16]) pytest.mark.parametrize(p, [0.0, 1.0]) def test_convention_mix_labels_stay_exact_for_half_precision_images(self, image_dtype, p): # Fixed by #4661: labels used to be cast to the image dtype, so bfloat16 returned [256, 1000] here (#4657). image torch.rand(2, 1, 4, 4, dtypeimage_dtype) labels torch.tensor([257, 999]) ... assert output.dtype image_dtype and mixed.dtype torch.float32 assert mixed[..., 0].flatten().tolist() [257.0, 999.0]它同时断言三件事图像输出 dtype 与输入一致output.dtype image_dtype标签输出为 float32mixed.dtype torch.float32大类别 id 精确保留257.0、999.0不丢精度。此外回放replay路径也有配套测试test_convention_mix_forward_parameters_replay_with_class_4706验证forward_parameters()产出的字典不含dtype键forward从输入张量实时取 dtype因而在 float32 上采样、float64 上回放时标签会正确输出 float64而非沿用旧字典test_conventions_mix.py#L560-L597。实战用法与注意事项以下示例可直接运行展示 bfloat16 输入下的 dtype 保持行为import torch import kornia as K torch.manual_seed(1) input torch.rand(2, 1, 3, 3, dtypetorch.bfloat16) # 半精度图像 label torch.tensor([257, 999]) # 超过半精度表示范围的大类别 id mixup K.RandomMixUpV2(data_keys[input, class]) images, labels mixup(input, label) assert images.dtype torch.bfloat16 # 图像保持在 bfloat16 assert labels.dtype torch.float32 # 标签提升为 float32id 精确 print(labels) # 形如 tensor([[257., 999., 0.19], ...])CutMix 用法一致cutmix K.RandomCutMixV2(data_keys[input, class], use_correct_lambdaTrue) images, labels cutmix(input, label) # labels 形状为 (num_mix, B, 3)配合 mix 标签训练时损失计算按 docstring 中的约定写法mixup.py#L46-L58import torch.nn.functional as F def loss_mixup(y, logits): criterion F.cross_entropy loss_a criterion(logits, y[:, 0].long(), reductionnone) loss_b criterion(logits, y[:, 1].long(), reductionnone) return ((1 - y[:, 2]) * loss_a y[:, 2] * loss_b).mean()使用中的关键注意事项data_keys[input, class]是启用标签混合的前提标签必须是 1 维张量shape(B,)输出为浮点标签矩阵p是批次级门控每次调用只采样一次选中则整批混合未选中p0时图像原样返回、标签列重复原始标签两次且 lambda 为 0见 mixup.py#L78-L91 的 Convention 说明半精度输入下不要期望标签也是半精度——这是有意为之的精度保护标签保持 float32 不会破坏与图像 dtype 的配对使用整数张量输入仍会被拒绝混合前请先归一化到浮点域uint8图像需先除以 255 转换回放采样参数params时dtype 永远以当前输入张量为准不要依赖旧参数字典中记录的 dtype。小结changelog.d/4471.fixed.md记录的修复虽然只有一句话却在源码层面涉及三条关键链路MixAugmentationBaseV2.transform_tensor的 dtype 白名单bfloat16 纳入、MixUp/CutMix 标签分支的calc_dtype提升规则半精度 → float32、以及回放路径的输入实时取 dtype 逻辑。配合tests/augmentation/test_conventions_mix.py中针对257/999这类大类别 id 的回归测试Kornia 的 mix 增强现在可以在 AMPAutomatic Mixed Precision训练流程中安全使用图像保持半精度以节省显存标签提升为 float32 以保证分类精度。【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/gh_mirrors/ko/kornia创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考