• 计算机视觉
• 深度学习
• 图像处理

本文基于 changelog.d/+migration-122.fixed.md 记录的迁移修复（issue #4043）展开，聚焦 Kornia 中rgb_to_yuv、yuv_to_rgb、rgb_to_xyz、xyz_to_rgb四个颜色空间转换函数在旧版 PyTorch 下无法通过torch.jit.script编译的问题。文章从根因（PEP 604 类型注解与 TorchScript 编译器的兼容性）、修复方案（回退到Optional[torch.Tensor]）、底层实现（共享的线性变换助手）与测试保障四个层面展开，读者将掌握在 Kornia 这类同时支持 TorchScript /torch.compile/ ONNX 导出的库中如何正确书写函数注解，以及如何在本地复现与验证该类问题。

问题背景：四个颜色转换函数在旧版 PyTorch 上无法 Script 化

Kornia 的kornia.color模块提供了一系列基于张量的颜色空间转换函数。本次修复涉及其中四个常用的线性转换：

• rgb_to_yuv/yuv_to_rgb：RGB 与 YUV（M/PAL 制式，见 ITU-R BT.470-5）之间的转换，定义于 kornia/color/yuv.py；
• rgb_to_xyz/xyz_to_rgb：RGB 与 CIE XYZ（D65 白点）之间的转换，定义于 kornia/color/xyz.py。

它们的特点是都通过 3×3 线性变换矩阵作用于通道维，并且共享同一个底层助手函数kornia.color.utils._apply_linear_transformation（见 kornia/color/utils.py）。正是这个共享函数的签名注解，成为四个转换函数在 TorchScript 编译链路上的共同瓶颈。

迁移记录明确指出：在 PyTorch 2.1.2 上对这四个函数调用torch.jit.script会直接失败，而在 PyTorch 2.9.1 上则能够正常编译。这意味着问题既不是算法实现错误，也不是设备或数据类型相关的运行时错误，而是编译期（scripting 阶段）对类型注解的解析能力差异——这正是 TorchScript 开发者最容易踩的"软地板"。

根因剖析：PEP 604 联合类型与 TorchScript 编译器的兼容性边界

什么是 PEP 604 联合类型

从 Python 3.10 起，PEP 604 允许用X | Y语法表达联合类型，例如torch.Tensor | None。这种写法比传统的Optional[torch.Tensor]更简洁，因此在新代码中非常流行。

为什么 TorchScript 会拒绝它

修复前的_apply_linear_transformation第三个参数bias的可选注解使用了 PEP 604 形式：

def _apply_linear_transformation( image: torch.Tensor, kernel: torch.Tensor, bias: torch.Tensor | None = None, # 修复前：TorchScript 编译器（旧版）不接受 ) -> torch.Tensor:

TorchScript 编译器在 script 模式下需要把 Python 函数完整地静态解析为 TorchScript IR。对于声明的最低支持 PyTorch 版本（本项目pyproject.toml中声明torch>=2.5.1，但迁移记录在 PyTorch 2.1.2 上即可复现），其编译器尚不能理解 PEP 604 的|联合类型语法，于是对这四个转换函数中任意一个执行torch.jit.script都会抛错——因为它们在函数体内都调用了这个带"问题注解"的助手。

值得注意的是，迁移记录给出的复现对比极具诊断价值：PyTorch 2.1.2 编译失败，而 PyTorch 2.9.1 编译成功。这说明新版编译器已补齐了对该语法的支持；但对于一个需要声明较宽 PyTorch 版本兼容范围的库来说，不能依赖"用户一定会升级到新版 PyTorch"这一假设，而必须让代码在声明支持的全部版本上都能编译。

仓库中的旁证：Ruff 规则 UP007

这一兼容性约束并非本次修复才引入的临时判断。在项目的 pyproject.toml 中，Ruff 的 lint 配置明确忽略了两条与注解风格相关的规则，并注明了原因：

"UP007", # Prefer Optional[], Union[] over | due to torch jit scripting "UP006", # Prefer List[], over list due to torch jit scripting

也就是说，Kornia 的编码规范本身就刻意避免在库代码中使用 PEP 604 的|和内置泛型list[...]，以保证 TorchScript 编译器在低版本 PyTorch 上也能解析。本次 #4043 修复正是这一规范在颜色模块上的落地执行。

修复方案：统一回退到Optional[torch.Tensor]

迁移记录的修复动作非常明确：将bias参数的类型注解从 PEP 604 形式改回Optional[torch.Tensor]。

修复后的签名（与当前仓库 kornia/color/utils.py 一致）：

from typing import Optional def _apply_linear_transformation( image: torch.Tensor, kernel: torch.Tensor, bias: Optional[torch.Tensor] = None ) -> torch.Tensor:

Optional[torch.Tensor]是typing模块自 Python 3.5 起就提供的标准形式，任何受支持的 PyTorch 版本的 TorchScript 编译器都能正确解析，因此一处修改即可同时修复四个函数的 script 化问题。

修复后的使用方式没有任何变化，四个转换函数依然像从前一样工作。例如 kornia/color/yuv.py 中rgb_to_yuv的调用：

KORNIA_CHECK_SHAPE(image, ["*", "3", "H", "W"]) dtype = image.dtype if image.is_floating_point() else torch.float32 kernel = torch.tensor( [ [0.299, 0.587, 0.114], [-0.147, -0.289, 0.436], [0.615, -0.515, -0.100], ], device=image.device, dtype=dtype, ) return _apply_linear_transformation(image, kernel)

rgb_to_xyz的调用方式相同（kornia/color/xyz.py），只是替换为 CIE RGB→XYZ 的 D65 白点矩阵。

底层实现：_apply_linear_transformation究竟做了什么

理解这次修复，还需要知道被修复的助手函数本身的行为，因为它决定了为何四个函数会共享同一个实现，也决定了修复的波及面。

从 kornia/color/utils.py 可以看到：

1. 类型与设备适配：若输入为整数张量则先提升为float32计算；kernel 与 bias 的 dtype/device 会向输入对齐（可传播float64精度）。
2. CPU 分支（含空张量）：使用torch.einsum("oi, ...ihw -> ...ohw", kernel, image)完成 3×3 通道线性变换，bias通过view(-1, 1, 1)广播后相加。源码注释说明，实测表明 einsum 在 CPU 上对该模式更快，且能正确处理所有空 shape，同时保持image、kernel、bias都在计算图中。
3. GPU/加速器分支：将输入 reshape 为(-1, 3, H, W)后调用F.conv2d，kernel 视为3×3×1×1的卷积权重，bias直接作为卷积偏置；输出再 reshape 回原形状。注释说明 conv2d 在 GPU/CUDA 上可获得显著加速。
4. 设备相关的分支选择：image_compute.device.type == "cpu" or image_compute.numel() == 0走 einsum，否则走 conv2d。

也就是说，四个颜色转换函数只是各自提供 3×3 系数矩阵与可选偏置，真正的张量运算全部收敛到这个共享助手。这解释了为什么一个注解问题会让四个函数同时中招——它们是同一段 TorchScript 编译路径上的不同入口。

此外，YUV 方向的实现还体现了精度设计：yuv_to_rgb的逆变换 kernel 并非独立四舍五入的公开逆公式，而是rgb_to_yuvkernel 的精确有理逆（注释中给出了精确分数，如[[1, -1/25344, 144439/126720], ...]），从而保证 RGB → YUV → RGB 往返误差仅受输入 dtype 精度限制（见 kornia/color/yuv.py）。这一细节与 #4044 的往返精度测试相互印证。

测试保障：test_jit如何防止回归

Kornia 的每个颜色转换函数都配套了 JIT 脚本化测试，这是本次修复能被稳定验证的基石。以 tests/color/test_yuv.py 中的TestRgbToYuv::test_jit为例：

def test_jit(self, device, dtype): img = torch.ones(2, 3, 4, 4, device=device, dtype=dtype) op = kornia.color.rgb_to_yuv op_jit = torch.jit.script(op) self.assert_close(op(img), op_jit(img))

测试同时覆盖多组(device, dtype)参数化组合：只要任意受支持版本在 script 阶段抛错，测试套件就会立即暴露问题；而assert_close则进一步保证 script 化后的输出与 eager 模式逐位一致。

类似的测试还出现在：

• tests/color/test_yuv.py 附近YuvToRgb的test_jit；
• tests/color/test_xyz.py 中RgbToXyz的test_jit以及同文件后续的XyzToRgb对应测试。

值得留意的是，这些测试文件中还包含test_dynamo（覆盖torch.compile/torch_optimizer路径）与test_round_trip_rgb_yuv_rgb（覆盖往返精度，见 tests/color/test_yuv.py）。这反映出 Kornia 对颜色转换模块的兼容性测试是多维度的：eager、JIT script、dynamo、ONNX 导出（ONNX_DEFAULT_INPUTSHAPE类属性）层层把关，任何一处注解或实现改动都可能牵动多条执行路径。

复现与验证：如何在本仓库中确认修复

如果你希望在本仓库中亲自复现"修复前失败、修复后通过"的过程，可以按如下步骤：

1. 在 Python 环境中安装当前仓库（requires-python = ">=3.11"，依赖torch>=2.5.1，见 pyproject.toml），并准备一个旧版 PyTorch 环境（如 2.1.x，用于复现失败场景）。
2. 对四个函数分别执行脚本化：

import torch import kornia for op in (kornia.color.rgb_to_yuv, kornia.color.yuv_to_rgb, kornia.color.rgb_to_xyz, kornia.color.xyz_to_rgb): scripted = torch.jit.script(op) # 旧版 PyTorch 上修复前会抛错 img = torch.rand(2, 3, 8, 8) print(scripted(img).shape) # 修复后 eager 与 script 结果一致

3. 运行颜色模块的既有测试验证无回归：

pytest tests/color/test_yuv.py tests/color/test_xyz.py

在修复后的代码上，四个函数均能顺利生成ScriptFunction，输出 shape 保持(*, 3, H, W)不变。

开发者启示：Kornia 的 TorchScript 兼容性实践

这次修复虽小，却浓缩了 Kornia 在多后端支持上的几条可复用经验：

1. 库代码的注解要面向最保守的执行环境。|联合类型和内置泛型虽然 Python 3.10+ 都支持，但 TorchScript 编译器对它们的接受度随 PyTorch 版本而异。Kornia 的做法是：库代码统一使用Optional[...]/List[...]/Tuple[...]，并在 Ruff 配置中显式禁用 UP006/UP007 以强制这一约定（pyproject.toml）。
2. 共享助手函数是兼容性风险汇聚点。_apply_linear_transformation一处注解错误，会让所有调用方在编译期集体失败。对这类被广泛调用的底层函数，签名必须格外保守。
3. 用测试矩阵锁定编译期行为。test_jit不是可选项：只要声明支持某个 PyTorch 版本下限，就必须让 CI 在该下限版本上真正跑一遍 script 化测试，否则"旧版编译失败"这类问题只能在用户报障时才发现。
4. 修复应保持行为不变。本次修复只改注解，不改算法、不改 API 签名、不改返回值语义，因此四个函数的 eager 结果、往返精度与梯度行为完全不受影响——这也是一个高质量兼容性修复应有的样子。

如果你正在自己的项目中编写会被torch.jit.script或torch.compile消费的代码，这条迁移记录是一个很好的反面教材：先写Optional[torch.Tensor]，不要贪图 PEP 604 的简洁；等确认所有目标 PyTorch 版本的编译器都支持后，再考虑是否迁移到新语法。

• 计算机视觉
• 深度学习
• 图像处理