Kornia 增强模块 CUDAtorch.compile常量搬运规避与采样器设备/dtype 迁移修复【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/gh_mirrors/ko/kornia本篇文章围绕 Kornia 增强augmentation模块在 CUDA 与torch.compile/make_fx编译场景下的两个核心问题展开一是将「常量张量在编译图中被搬运到主机host再拷贝回设备」的规避方案扩展到 shear、thin-plate spline、erasing、crop、resize、3D affine/perspective 及光照/混合参数生成器等更多算子二是修复RandomShear、RandomAffine采样器迁移及线性光照生成器混合设备符号生成等设备/dtype 一致性缺陷。读完本文你将理解 Kornia 如何在 Eager、Dynamo 与make_fx三种执行模式下用统一的无主机拷贝标量工厂构建常量以及如何通过set_rng_device_and_dtype保证采样参数落在正确的设备与 dtype 上。背景为什么编译图中会出现「常量搬运」在使用 CUDA 的torch.compile场景中InductorPyTorch 的编译后端可以复用被 lifted 到图外的 CPU 常量张量而无需将其传输到 CUDA kernel 中对应 PyTorch issue pytorch/pytorch#196969。这种隐式的主机-设备host-device拷贝会带来性能损失并且在某些图捕获路径下会破坏常量语义。Kornia 增强模块的做法是在图内in-graph用**标量工厂scalar factory和栈stack**构造小常量完全绕开「把 CPU 张量 lifted 到图外」的路径包括那些与 CUDA 张量一起返回的常量。与此同时索引式标量写入indexed scalar writes同样是不安全的trace 时会 lift 它们的右值。统一常量构造工具_constant_tensorKornia 将上述策略收敛为一个内部工具函数_constant_tensor实现在 kornia/augmentation/utils/helpers.pydef _constant_tensor( data: Union[float, List[Any], Tuple[Any, ...]], *, device: Union[str, torch.device, None] None, dtype: torch.dtype, ) - torch.Tensor:其设计要点包括数据形态data只允许 Python 标量或矩形的嵌套 list/tuple不允许张量其他 array-like如 NumPy 数组保持torch.as_tensor语义。标量分支对int、float、torch.SymInt、torch.SymFloat直接调用torch.full((), data, devicedevice, dtypedtype)在指定设备上就地构造。嵌套结构分支通过_flatten_constant递归展开出形状再为每个叶子值填充一次最后torch.stack还原形状。去重优化每个不同的 Python 标量只填充一次单个stack复用——因此像 box 角点这种重复坐标每个值只触发一个 kernel。符号尺寸处理符号尺寸symbolic sizes不合并因为比较它们会增加 guardsfloat key 携带符号因为-0.0 0.0NaN 则每个叶子单独填充因为NaN ! NaN。注释中明确说明Eager 执行、Dynamotorch.compile与make_fx都运行同一套构造逻辑全程不产生 host-device 拷贝。调用方需要显式选择 dtype并在需要「先取整再转类型」时在调用该工具之前完成坐标运算。增强基类中的应用在增强模块基类 kornia/augmentation/base.py 中forward_parameters使用_constant_tensor构造forward_input_shapebase.py#L293-L303def forward_parameters(self, batch_shape: Tuple[int, ...]) - Dict[str, torch.Tensor]: batch_prob self.__batch_prob_generator__(batch_shape, self.p, self.p_batch, self.same_on_batch) _params self.generate_parameters(batch_shape) if _params is None: _params {} _params[batch_prob] batch_prob # Added another input_size parameter for geometric transformations # This might be needed for correctly inversing. input_size _constant_tensor(batch_shape, dtypetorch.long) _params.update({forward_input_shape: input_size}) return _params同时基类的__batch_prob_generator__base.py#L210-L244也贯彻了「避免 graph break」的编译友好理念p 1/p 0/same_on_batch这些分支基于Python 值判断因此在 trace 时即可解析不会产生 graph break原先依赖数据的if batch_prob.sum() 1分支被替换为无分支的batch_prob batch_prob * elem_prob当 batch 被选中时结果为elem_prob未被选中时全为 0与旧分支行为完全一致但不再触发 graph break。基类中另一处与编译/导出相关的处理是_commit_statebase.py#L246-L269在torch.export捕获期间会跳过属性写入is_exporting()时直接返回因为导出的图拒绝forward中的属性变更被捕获的图像输出不受影响只是事后读取状态在导出图中无意义被跳过。迁移到更多算子五类参数生成器的覆盖本次变更将上述常量搬运规避从既有算子扩展到以下五类参数生成器分布在 kornia/augmentation/random_generator/ 下几何形变类shear_2d/shear.py、thin-plate splineTPS、3D affine / 3D perspective_3d/affine.py、_3d/perspective.py裁剪/缩放类erasing、crop、resize_2d/crop.py、_2d/resize.py等光照与混合类illumination / mix 参数生成器gaussian_illumination.py、linear_illumination.py、cutmix.py、mosaic.py等。以 2D shear 生成器 kornia/augmentation/random_generator/_2d/shear.py 为例其forward中裁剪中心使用_constant_tensor构造center: torch.Tensor _constant_tensor([width, height], device_device, dtype_dtype).view(1, 2) / 2.0 - 0.5 center center.expand(batch_size, -1)注意这里的用法是先构造再运算/ 2.0 - 0.5即「先转类型再做坐标算术」的正确顺序。采样得到的sx、sy随后被to(device_device, dtype_dtype)归一。3D crop 生成器则是更典型的应用场景kornia/augmentation/random_generator/_3d/crop.py#L232-L267源/目标八个角点top-left-front、top-right-front、bottom-right-front、bottom-left-front、top-left-back、top-right-back、bottom-right-back、bottom-left-back都以嵌套 list 形式传入_constant_tensor(..., devicedevice, dtypetorch.long)再expand(batch_size, -1, -1)到 batch 维度——同一组角点值只填充一次避免每个 batch 元素重复触发 kernel。一个重要的边界Crop 类算子保留其既有的 graph breaksCrop operations retain their existing graph breaks。这属于有意为之的取舍——裁剪路径中的数据依赖分支没有被强行改写因此常量规避不适用于这些路径。RandomShear/RandomAffine采样器迁移修复问题现象修复前RandomShear和RandomAffine的采样器迁移存在缺陷标量或二元组形式的shear范围没有遵循请求的 device 与 dtype。典型后果是——把带标量或两个元素shear参数的RandomAffine迁移到 CUDA 时会抛出device-mismatch 错误。根因与修复根因在于采样器构建时对shear边界bound的推导没有走set_rng_device_and_dtype请求的设备。修复后ShearGenerator.make_samplers中_shear_bound(self.shear, device, dtype)生成的上下界、以及UniformDistribution采样器都落在正确的设备与 dtype 上kornia/augmentation/random_generator/_2d/shear.py#L70-L82def make_samplers(self, device: torch.device, dtype: torch.dtype) - None: _shear _shear_bound(self.shear, device, dtype) _joint_range_check(_shear[0], shear) _joint_range_check(_shear[1], shear) self.shear_x _shear[0].clone() self.shear_y _shear[1].clone() shear_x_sampler UniformDistribution(_shear[0][0], _shear[0][1], validate_argsFalse) shear_y_sampler UniformDistribution(_shear[1][0], _shear[1][1], validate_argsFalse)这里shear支持四种形态shear.py#L39-L47shear参数形态语义float沿 x 轴在(-shear, shear)范围内剪切(a, b)沿 x 轴在(-shear, shear)范围内剪切(a, b, c, d)x 轴剪切取(shear[0], shear[1])y 轴剪切取(shear[2], shear[3])torch.Tensor2x2x 轴剪切取(shear[0][0], shear[0][1])y 轴剪切取(shear[1][0], shear[1][1])无论哪种形态采样参数都会返回形状为(B,)的shear_x与shear_y。dtype 迁移的完整机制由基类提供set_rng_device_and_dtype(device, dtype)会同时更新 gate 与参数生成器的采样器base.py#L190-L208。需要留意其文档注释中的告诫返回的参数可能被转换到其他设备/dtype仅查看_params无法判断采样实际发生在哪里部分生成器仍可能保留内部 CPU 张量或忽略请求的精度个别生成器/设备组合在 forward 时仍可能失败——这些问题被跟踪在 issue #4426并在 get-started/conventions 页面说明。测试佐证仓库测试覆盖了该机制如 tests/augmentation/test_augmentation.py#L6246 通过aug.set_rng_device_and_dtype(devicedevice, dtypedtype)验证不同设备/dtype 组合tests/augmentation/test_base.py#L968-L1004 则对 CPU、CUDA、MPS 上的 RNG 状态与迁移行为做断言相关注释指出跨设备回归主要在 CUDA 与 MPS 上验证。线性光照生成器的混合设备符号生成修复第二个缺陷位于两个线性光照生成器将它们的采样器迁移到新设备后符号sign生成存在混合设备问题。以 2D 线性光照生成器 kornia/augmentation/random_generator/_2d/linear_illumination.py 为例修复后的make_samplers将所有采样边界显式.to(device, dtype)到目标设备linear_illumination.py#L58-L75def make_samplers(self, device: torch.device, dtype: torch.dtype) - None: gain _range_bound(self.gain, gain).to(device, dtype) self.gain_sampler UniformDistribution(gain[0], gain[1], validate_argsFalse) sign _range_bound(self.sign, sign, bounds(-1.0, 1.0), center0.0).to(device, dtype) self.sign_sampler UniformDistribution(sign[0], sign[1], validate_argsFalse) # Draw the directions on the sampler device but always in float32: on MPS, half-precision # torch.rand can return exactly 1.0, which would truncate to the invalid direction 4. self.directions_sampler UniformDistribution( torch.tensor(0.0, devicedevice, dtypetorch.float32), torch.tensor(4.0, devicedevice, dtypetorch.float32), validate_argsFalse, )forward中符号生成现在完全在同一设备上完成linear_illumination.py#L83-L93# Random gain and sign gain_factor _adapted_rsampling((batch_size, 1, 1, 1), self.gain_sampler, same_on_batch).to( device_device, dtype_dtype ) sign_positive _adapted_rsampling((batch_size, 1, 1, 1), self.sign_sampler, same_on_batch) 0.0 sign sign_positive.to(device_device, dtype_dtype) * 2 - 1符号被编码为sign sign_positive.to(...) * 2 - 1正采样得 1否则得 -1全程不再出现「采样器在 A 设备、符号张量在 B 设备」的混合设备局面。方向采样器刻意固定在 float32——这是 MPS 上规避半精度torch.rand恰好返回 1.0、进而截断成非法方向 4 的针对性设计。CenterCrop3D空 batch 参数 dtype/设备一致性最后一个修复针对 3D 中心裁剪的边界情况当 batch 为空时CenterCrop3D的参数此前是 CPU 上的float32张量与非空 batch请求设备上的long张量不一致。修复后空 batch 的参数同样成为请求设备上的long张量与非空 batch 行为对齐。这一改动与 3D crop 生成器整体改用_constant_tensor(..., dtypetorch.long)构造角点坐标的路径一致kornia/augmentation/random_generator/_3d/crop.py#L232-L267保证无论 batch 是否为空points_src/points_dst的类型与设备语义都稳定避免下游变换在空 batch 推理如导出、批处理打包时产生 dtype 分支。总结改动清单与验证方式本次修复可归纳为三个层面编译性能层把 CUDAtorch.compile下的常量搬运规避_constant_tensor 无分支 batch gate扩展到 shear、thin-plate spline、erasing、crop、resize、3D affine/perspective 与 illumination/mix 参数生成器Eager、Dynamo、make_fx三种执行模式统一使用无 host-device 拷贝的标量工厂构建常量crop 保留既有 graph breaks。设备/dtype 正确性层RandomShear/RandomAffine的标量与二元组shear范围遵循set_rng_device_and_dtype请求的设备与 dtype消除 CUDA 迁移时的 device-mismatch两个线性光照生成器的符号生成不再混合设备。边界一致性层CenterCrop3D空 batch 参数改为请求设备上的long张量与非空 batch 对齐。相关改动涉及的核心文件路径常量构造工具kornia/augmentation/utils/helpers.py增强基类batch gate、forward_input_shape、state 提交kornia/augmentation/base.py2D shear 生成器kornia/augmentation/random_generator/_2d/shear.py3D crop 生成器kornia/augmentation/random_generator/_3d/crop.py光照生成器kornia/augmentation/random_generator/_2d/linear_illumination.py、kornia/augmentation/random_generator/_2d/gaussian_illumination.py相关测试tests/augmentation/test_augmentation.py、tests/augmentation/test_base.py、tests/augmentation/test_augmentation_compile.py如果你正在使用 CUDA 上的torch.compile训练流水线建议重点回归验证RandomShear、RandomAffine尤其标量/二元组shear、3D crop 与光照类增强在编译前后、空 batch 与非空 batch 下的输出一致性。【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/gh_mirrors/ko/kornia创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
` 溢出修复)