企业官网建设流程全解析

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场景中，Inductor（PyTorch 的编译后端）可以复用被 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_tensor

Kornia 将上述策略收敛为一个内部工具函数_constant_tensor，实现在 kornia/augmentation/utils/helpers.py：

def _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, device=device, dtype=dtype)，在指定设备上就地构造。
• 嵌套结构分支：通过_flatten_constant递归展开出形状，再为每个叶子值填充一次，最后torch.stack还原形状。
• 去重优化：每个不同的 Python 标量只填充一次，单个stack复用——因此像 box 角点这种重复坐标，每个值只触发一个 kernel。
• 符号尺寸处理：符号尺寸（symbolic sizes）不合并，因为比较它们会增加 guards；float key 携带符号（因为-0.0 == 0.0），NaN 则每个叶子单独填充（因为NaN != NaN）。

注释中明确说明：Eager 执行、Dynamo（torch.compile）与make_fx都运行同一套构造逻辑，全程不产生 host-device 拷贝。调用方需要显式选择 dtype，并在需要「先取整再转类型」时，在调用该工具之前完成坐标运算。

增强基类中的应用

在增强模块基类 kornia/augmentation/base.py 中，forward_parameters使用_constant_tensor构造forward_input_shape（base.py#L293-L303）：

def 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, dtype=torch.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_state（base.py#L246-L269）：在torch.export捕获期间会跳过属性写入（is_exporting()时直接返回），因为导出的图拒绝forward中的属性变更；被捕获的图像输出不受影响，只是事后读取状态（在导出图中无意义）被跳过。

迁移到更多算子：五类参数生成器的覆盖

本次变更将上述常量搬运规避从既有算子扩展到以下五类参数生成器（分布在 kornia/augmentation/random_generator/ 下）：

1. 几何形变类：shear（_2d/shear.py）、thin-plate spline（TPS）、3D affine / 3D perspective（_3d/affine.py、_3d/perspective.py）；
2. 裁剪/缩放类：erasing、crop、resize（_2d/crop.py、_2d/resize.py等）；
3. 光照与混合类：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(..., device=device, dtype=torch.long)，再expand(batch_size, -1, -1)到 batch 维度——同一组角点值只填充一次，避免每个 batch 元素重复触发 kernel。

一个重要的边界：Crop 类算子保留其既有的 graph breaks（Crop 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-L82）：

def 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_args=False) shear_y_sampler = UniformDistribution(_shear[1][0], _shear[1][1], validate_args=False)

这里shear支持四种形态（shear.py#L39-L47）：

无论哪种形态，采样参数都会返回形状为(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(device=device, dtype=dtype)验证不同设备/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-L75）：

def 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_args=False) sign = _range_bound(self.sign, "sign", bounds=(-1.0, 1.0), center=0.0).to(device, dtype) self.sign_sampler = UniformDistribution(sign[0], sign[1], validate_args=False) # 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, device=device, dtype=torch.float32), torch.tensor(4.0, device=device, dtype=torch.float32), validate_args=False, )

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(..., dtype=torch.long)构造角点坐标的路径一致（kornia/augmentation/random_generator/_3d/crop.py#L232-L267），保证无论 batch 是否为空，points_src/points_dst的类型与设备语义都稳定，避免下游变换在空 batch 推理（如导出、批处理打包）时产生 dtype 分支。

总结：改动清单与验证方式

本次修复可归纳为三个层面：

1. 编译性能层：把 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。
2. 设备/dtype 正确性层：RandomShear/RandomAffine的标量与二元组shear范围遵循set_rng_device_and_dtype请求的设备与 dtype，消除 CUDA 迁移时的 device-mismatch；两个线性光照生成器的符号生成不再混合设备。
3. 边界一致性层：CenterCrop3D空 batch 参数改为请求设备上的long张量，与非空 batch 对齐。

相关改动涉及的核心文件路径：

• 常量构造工具：kornia/augmentation/utils/helpers.py
• 增强基类（batch gate、forward_input_shape、state 提交）：kornia/augmentation/base.py
• 2D shear 生成器：kornia/augmentation/random_generator/_2d/shear.py
• 3D 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