Kornia 修复 MPS/CUDA 跨设备增强 RuntimeError：batch_prob 掩码的设备迁移与 branchless 混合原理

计算机视觉人工智能深度学习图像处理【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址https://gitcode.com/gh_mirrors/ko/kornia点击查看免费下载本指南讲解 Kornia 增强模块在 CPU 生成随机参数、输入位于 MPS 或 CUDA 加速器时抛出跨设备RuntimeError的根因与官方修复方案changelog 条目migration-109.fixed.mdissue #4164。读完本文你将掌握batch_prob门控掩码的生成与消费路径、四处torch.wherebranchless 混合的实现细节以及修复后augmentation(x_accelerator)的正确用法与残留限制。问题背景默认 CPU 随机数生成与加速器输入之间的矛盾Kornia 的增强模块kornia.augmentation在设计上有一个重要约定随机参数默认在 CPU 上生成与输入张量的设备无关。这一约定体现在 kornia/augmentation/base.py 的基类构造函数中def __init__( self, p: float 0.5, p_batch: float 1.0, same_on_batch: bool False, keepdim: bool False, ) - None: super().__init__() self.p p self.p_batch p_batch self.same_on_batch same_on_batch self.keepdim keepdim self._params: Dict[str, torch.Tensor] {} self._param_generator: Optional[RandomGeneratorBase] None self.flags: Dict[str, Any] {} self.set_rng_device_and_dtype(torch.device(cpu), torch.get_default_dtype())也就是说模块实例化后其 RNG 相关状态包括self.device与self.dtype默认落在 CPU。这样做的好处是跨设备可复现同一随机种子在不同设备上采样出完全一致的增强参数便于训练/评估结果对齐。其中最关键的门控张量是batch_prob它由__batch_prob_generator__生成kornia/augmentation/base.pydef __batch_prob_generator__( self, batch_shape: Tuple[int, ...], p: float, p_batch: float, same_on_batch: bool, ) - torch.Tensor: batch_prob: torch.Tensor if p_batch 1: batch_prob torch.ones(1, deviceself.device, dtypeself.dtype) elif p_batch 0: batch_prob torch.zeros(1, deviceself.device, dtypeself.dtype) else: batch_prob (torch.rand(1, deviceself.device) p_batch).to(self.dtype) elem_prob: torch.Tensor if p 1: elem_prob torch.ones(batch_shape[0], deviceself.device, dtypeself.dtype) elif p 0: elem_prob torch.zeros(batch_shape[0], deviceself.device, dtypeself.dtype) elif same_on_batch: elem_prob (torch.rand(1, deviceself.device) p).to(self.dtype).expand(batch_shape[0]) else: elem_prob (torch.rand(batch_shape[0], deviceself.device) p).to(self.dtype) # Branchless combine (replaces the>staticmethod def _blend_by_prob( transformed: torch.Tensor, not_transformed: torch.Tensor, to_apply: torch.Tensor ) - torch.Tensor: Select transformed vs non-transformed samples element-wise by to_apply. When the two branches share a shape this is a torch.where blend (onnx- and fullgraph-friendly). Shape-changing augmentations (e.g. crop/resize) whose branches differ in spatial size fall back to a Python branch on to_apply.any(), which is not onnx-exportable. if transformed.shape not_transformed.shape and transformed.shape[0] to_apply.shape[0]: to_apply_expanded to_apply.view(-1, *([1] * (len(transformed.shape) - 1))).to(transformed.device) return torch.where(to_apply_expanded, transformed, not_transformed) return transformed if bool(to_apply.any()) else not_transformed关键一行是to_apply_expanded to_apply.view(-1, *([1] * (len(transformed.shape) - 1))).to(transformed.device)先沿 batch 维把形状为(B,)的掩码广播展开为(B, 1, 1, ...)再迁移到transformed所在设备最后执行torch.where。2. 2D 几何/仿射矩阵路径kornia/augmentation/_2d/base.py 的RigidAffineAugmentationBase2D.generate_transformation_matrixbatch_prob params[batch_prob] to_apply torch.atleast_1d(batch_prob 0.5) in_tensor self.transform_tensor(input) trans_matrix_applied self.compute_transformation(in_tensor, paramsparams, flagsflags) if self.p 1.0 and self.p_batch 1.0: # Always applied (static probabilities): skip building the identity and the where. trans_matrix trans_matrix_applied if is_autocast_enabled(): trans_matrix trans_matrix.type(input.dtype) return trans_matrix trans_matrix_identity self.identity_matrix(in_tensor) if is_autocast_enabled(): trans_matrix_applied trans_matrix_applied.type(input.dtype) trans_matrix_identity trans_matrix_identity.type(input.dtype) if trans_matrix_applied.shape[0] to_apply.shape[0] trans_matrix_identity.shape[0]: to_apply_expanded to_apply.view(-1, *([1] * (trans_matrix_applied.dim() - 1))).to( trans_matrix_applied.device ) trans_matrix torch.where(to_apply_expanded, trans_matrix_applied, trans_matrix_identity) else: # e.g. VideoSequential passes B-sized batch_prob into a B*T-sized input trans_matrix trans_matrix_applied if bool(to_apply.any()) else trans_matrix_identity这里还有一处值得注意的配套优化当p 1.0且p_batch 1.0静态全应用时直接返回计算矩阵跳过身份矩阵的构建与torch.where混合——注释指出该矩阵路径约占一次 flip forward 开销的 40%这是针对热路径的额外加速与本次设备修复同属本次变更的一部分。3. 3D 矩阵路径kornia/augmentation/_3d/base.py 与 2D 版本对称同样在torch.where前执行.to(trans_matrix_applied.device)batch_prob params[batch_prob] to_apply torch.atleast_1d(batch_prob 0.5) in_tensor self.transform_tensor(input) trans_matrix_applied self.compute_transformation(in_tensor, paramsparams, flagsflags) trans_matrix_identity self.identity_matrix(in_tensor) if trans_matrix_applied.shape[0] to_apply.shape[0] trans_matrix_identity.shape[0]: to_apply_expanded to_apply.view(-1, *([1] * (trans_matrix_applied.dim() - 1))).to( trans_matrix_applied.device ) trans_matrix torch.where(to_apply_expanded, trans_matrix_applied, trans_matrix_identity) else: trans_matrix trans_matrix_applied if bool(to_apply.any()) else trans_matrix_identity4. Mix 增强路径kornia/augmentation/_2d/mix/base.py 的transform_inputbatch_prob params[batch_prob] to_apply torch.atleast_1d(batch_prob 0.5) ori_shape input.shape in_tensor self.transform_tensor(input) # Compute the non-transform branch first; if no element is to be transformed, short-circuit # (mix transforms like RandomJigsaw subset their input internally and cant operate on an # empty subset). non_applied self.apply_non_transform(in_tensor, params, flags) if not bool(to_apply.any()): output non_applied return _transform_output_shape(output, ori_shape) if self.keepdim else output applied self.apply_transform(in_tensor, params, flags) applied_post self.apply_non_transform(applied, params, flags) if applied_post.shape non_applied.shape and applied_post.shape[0] to_apply.shape[0]: to_apply_expanded to_apply.view(-1, *([1] * (applied_post.dim() - 1))).to(applied_post.device) output torch.where(to_apply_expanded, applied_post, non_applied) else: # Shape-changing mix augmentations (e.g. RandomMosaic with different output_size) # cannot be where-blended. Fall back to the all-applied branch. output applied_post return _transform_output_shape(output, ori_shape) if self.keepdim else outputMix 路径除了设备迁移还保留了“形状变化时退化为全应用分支”的兜底逻辑——例如RandomMosaic在output_size与输入尺寸不一致时无法做逐元素where混合只能整体应用。修复后的行为与正确用法修复后最直观的变化是直接对加速器张量调用增强模块即可无需再手动迁移模块或其 RNG。修复前需要这样绕过#4151 中的 workaroundimport torch from kornia.augmentation import RandomHorizontalFlip x torch.randn(4, 3, 224, 224, devicecuda) # 或 mps # 修复前必须把模块连同其 RNG 状态先搬到加速器 aug RandomHorizontalFlip(p0.5).to(cuda) y aug(x)修复后可以直接调用import torch from kornia.augmentation import RandomHorizontalFlip x torch.randn(4, 3, 224, 224, devicecuda) # 或 torch.device(mps) aug RandomHorizontalFlip(p0.5) # 保持默认RNG 在 CPU y aug(x) # 修复后正常返回无跨设备错误这一用法对所有继承链均生效2D 强度/几何增强走_blend_by_prob与 2D 矩阵路径3D 增强走 3D 矩阵路径如 kornia/augmentation/_3d/base.py 的generate_transformation_matrixMix 增强走 kornia/augmentation/_2d/mix/base.py 的transform_input包括 MixUp、CutMix、RandomJigsaw、RandomMosaic 等。同时保持了两个既有优势不变CPU 随机数生成batch_prob与各类参数仍在 CPU 采样torch.manual_seed后在不同设备上得到一致的采样序列跨设备可复现性不受影响ONNX / fullgraph 友好混合仍是无 Python 分支的torch.where形式torch.compile(fullgraphTrue)与 ONNX 导出路径不受破坏形状变化的增强除外见下文。边界情况与已知限制修复并非万能以下几点需要在使用中留意均可在源码注释与文档中找到依据形状变化的增强仍不可 ONNX 导出_blend_by_prob的 docstring 明确说明当变换/非变换分支形状不同如 crop、resize 类时会退化为基于to_apply.any()的 Python 分支该路径不可 ONNX 导出torch.compile的fullgraphTrue也因此对这类增强不保证成立kornia/augmentation/base.py。静态概率快速路径p 1.0 and p_batch 1.0时直接返回变换结果连batch_prob门控都不计算kornia/augmentation/_2d/base.py这同时让Resize这类形状变化的增强在静态概率下可以 fullgraph 编译。mix 的 mask/boxes/keypoints/class 数据键不走torch.wherekornia/augmentation/_2d/mix/base.py 中的transform_mask、transform_boxes、transform_keypoint、transform_class使用sum(to_apply)的 Python 分支选择不涉及跨设备混合因此不受此 bug 影响。set_rng_device_and_dtype迁移不完整文档与源码kornia/augmentation/base.py均提醒该方法会更新门控与参数生成器的采样器但部分生成器仍保留内部 CPU 张量或忽略指定精度某些生成器/设备组合在 forward 期间仍可能失败issue #4426。因此推荐做法仍是“RNG 留在 CPU、由混合边界负责掩码迁移”而不是依赖set_rng_device_and_dtype把采样搬到加速器。B 0空批次空批次在多数类上返回空输出但并非全库保证少数类在B 0时会抛异常。测试验证本次修复的回归验证覆盖在增强测试套件中。tests/augmentation/container/test_augmentation_sequential.py通过device/dtype参数化测试包括使用get_cuda_or_mps_device_if_available这类测试辅助函数获取可用加速器其名称出现在 tests/api_surface.json 的公开工具清单中在 CUDA 与 MPS 设备上验证 2D、3D 与 mix 增强的前向行为其中包含对 MPS 上float64输入的专门判断if device.type mps and image_dtype torch.float64。结合本修复这些测试确认了augmentation(x_accelerator)在各设备上的正常执行。小结migration-109这项修复issue #4164解决的是一个典型的“设备归属”问题torch.where的掩码与操作数必须同设备而 Kornia 有意让随机参数留在 CPU。修复在四个 branchless 混合边界统一执行to_apply.to(transformed.device)以极小的改动同时满足了三点诉求——加速器输入直接可用、CPU RNG 跨设备可复现、ONNX/fullgraph 友好的混合结构不受破坏。理解这条变更的细节有助于你在使用 Kornia 增强管线时正确安排模块与数据的设备关系并规避形状变化增强与 RNG 迁移方面的已知坑位。赞分享计算机视觉人工智能深度学习图像处理【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址https://gitcode.com/gh_mirrors/ko/kornia点击查看免费下载相关推荐Kornia 跨设备增强修复解析CPU 随机掩码与 MPS/CUDA 输入的设备一致性方案Kornia 跨设备增强修复解析CPU 随机掩码与 MPS/CUDA 输入的设备一致性方案 本篇技术文章围绕 Kornia 仓库 changelog 中的 计算机视觉深度学习人工智能图像处理Kornia 增强模块 CUDA torch.compile 常量搬运规避与采样器设备/dtype 迁移修复Kornia 增强模块 CUDA torch.compile 常量搬运规避与采样器设备/dtype 迁移修复 本篇文章围绕 Kornia 增强augmenta计算机视觉人工智能深度学习图像处理Kornia RandomAffine 与 RandomPerspective 的 CUDA 编译修复与随机生成器设备迁移语义Kornia RandomAffine 与 RandomPerspective 的 CUDA 编译修复与随机生成器设备迁移语义 导读 本篇文章围绕 Kornia计算机视觉深度学习人工智能图像处理创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考