ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

Kornia 修复 MPS 等加速器后端的空张量转换:YUV420/YUV422 空输入语义与 reshape 歧义问题深度解析

Kornia 修复 MPS 等加速器后端的空张量转换:YUV420/YUV422 空输入语义与 reshape 歧义问题深度解析 计算机视觉深度学习人工智能图像处理【免费下载链接】kornia 空间人工智能的几何计算机视觉库项目地址https://gitcode.com/kornia/kornia点击查看免费下载导读本文围绕 Kornia 仓库changelog.d/migration-103.fixed.md记录的修复对应 issue #4185剖析 YUV420/YUV422 颜色转换在 MPS 等加速器后端处理**空张量empty tensor**时的一个隐蔽缺陷reshape无法从零元素中推断 batch 维度导致空输入转换在 MPS 上报错而无法返回文档承诺的空 RGB 输出。文章将结合 kornia/color/yuv.py、kornia/color/utils.py 与 tests/color/test_yuv.py 的源码与测试讲清问题根因、修复策略、形状校验规则以及如何复现与验证。一、这条 changelog 记录了什么changelog.d/migration-103.fixed.md全文记录了本次修复的核心内容Preserve empty color tensors on accelerator backends instead of askingreshapeto infer an ambiguous batch dimension. This lets the YUV420 and YUV422 empty-input conversions return their documented empty RGB output on MPS, matching CPU behavior (#4185).翻译过来本次变更做了三件事保留preserve空颜色张量在加速器后端accelerator backends如 MPS不再要求reshape去推断一个有歧义的 batch 维度空输入转换行为与 CPU 对齐YUV420 与 YUV422 的空输入转换在 MPS 上可以返回其文档中声明的空 RGB 输出与 CPU 行为保持一致matching CPU behavior并关联了 issue #4185。也就是说这是一个典型的设备一致性device parity修复同样的输入在 CPU 上能跑通、在 MPS 上报错的函数通过保留空张量绕过reshape的歧义推断使各后端行为统一。二、前置背景YUV420 / YUV422 转换的输入输出契约要理解这条修复先要弄清涉及的四个转换函数及其形状契约。它们全部位于 kornia/color/yuv.py并统一从 kornia/color/init.py 导出rgb_to_yuv420、rgb_to_yuv422、yuv420_to_rgb、yuv422_to_rgb及对应的RgbToYuv420、RgbToYuv422、Yuv420ToRgb、Yuv422ToRgb等nn.Module封装。2.1 正向转换RGB → YUV420 / YUV422rgb_to_yuv420(image)输入形状(*, 3, H, W)返回(y, uv)二元组y亮度平面形状(*, 1, H, W)uv色度平面形状(*, 2, H/2, W/2)即 4:2:0 在宽高两个方向各降采样一半。rgb_to_yuv422(image)输入形状(*, 3, H, W)返回(y, uv)二元组y亮度平面形状(*, 1, H, W)uv色度平面形状(*, 2, H, W/2)即 4:2:2 仅在宽度方向降采样一半。两个函数共同的硬性前置条件yuv.py 与 yuv.pyif len(image.shape) 2 or image.shape[-2] % 2 1 or image.shape[-1] % 2 1: raise ShapeError(fInput HW must be evenly divisible by 2. Got {image.shape})即输入宽高必须能被 2 整除否则抛出kornia.core.exceptions.ShapeError。所有数值都遵循 M/PALBT.470-5的 YUV 模型Y 亮度在(0, 1)U 在(-0.436, 0.436)V 在(-0.615, 0.615)。2.2 逆向转换YUV420 / YUV422 → RGByuv420_to_rgb(imagey, imageuv)imagey形状(*, 1, H, W)imageuv形状(*, 2, H/2, W/2)要求色度的宽高恰好是亮度的一半yuv.pyyuv422_to_rgb(imagey, imageuv)imagey形状(*, 1, H, W)imageuv形状(*, 2, H, W/2)要求色度高度与亮度一致、宽度为亮度一半yuv.py。逆向转换的实现思路是先升采样再走 4:4:4 通道YUV420 用repeat_interleave(2, dim-1).repeat_interleave(2, dim-2)做 2×2 复制YUV422 只做宽度方向的repeat_interleave(2, dim-1)拼接回 3 通道后交给yuv_to_rgb。需要特别留意的是yuv_to_rgbyuv.py它的内核是rgb_to_yuv内核的精确逆矩阵按有理数求逆后精确舍入到 float64而不是另抄一份舍入后的公开逆变换公式因此 RGB → YUV → RGB 的往返误差只受输入 dtype 精度限制。这个细节解释了测试里为什么会有往返一致性forth_and_back测试。三、问题根源加速器后端reshape推断 batch 维度的歧义颜色转换最终都要经过 kornia/color/utils.py 中的_apply_linear_transformation。这个辅助函数根据设备选择了两条性能路径if image_compute.device.type cpu or image_compute.numel() 0: out torch.einsum(oi, ...ihw - ...ohw, kernel_compute, image_compute) ... # BRANCH 2: GPU/Accelerators (Conv2d) input_flat image_compute.reshape(-1, 3, input_shape[-2], input_shape[-1]) weight kernel_compute.view(3, 3, 1, 1) out_flat F.conv2d(input_flat, weight, biasbias_compute) out out_flat.reshape(input_shape)分支 1einsum适用于 CPU 以及所有空张量注释明确写道 Einsum handles every empty shape and keeps image, kernel, and bias in the graph即 einsum 能天然处理任意空形状分支 2conv2d适用于 GPU/加速器后端需要先把任意前导维度(*, 3, H, W)压平成(B, 3, H, W)再调用F.conv2d。问题恰恰出在分支 2 的reshape(-1, 3, input_shape[-2], input_shape[-1])上当张量为空numel() 0时-1表示的 batch 维度无法被唯一确定。例如形状为(3, 0, 0)的空张量它既可以理解为batch1, C3, H0, W0也可以理解为batch0, C3, H任意, W任意——reshape无法从 0 个元素中反推出一个确定的 batch 值这在部分后端尤其是 MPS会直接报错导致空输入转换失败。从源码结构看修复前的分支条件只判断了device.type cpu因此空张量在 MPS 上会落入 conv2d 分支、触发上述歧义修复后把image_compute.numel() 0也纳入 einsum 分支空张量无论在哪块后端都走 einsum 路径从而保留空颜色张量、避免 reshape 推断歧义。同理rgb_to_yuv420内部做 2×2 盒式降采样时也涉及显式 reshapeyuv.pyuv yuvimage[..., 1:3, :, :] leading list(uv.shape[:-3]) batch 1 for size in leading: batch * size height, width uv.shape[-2], uv.shape[-1] uv F.avg_pool2d(uv.reshape([batch, 2, height, width]), 2) return yuvimage[..., :1, :, :], uv.reshape(leading [2, height // 2, width // 2])对空输入而言batch是前导维度之积可能为 0avg_pool2d前后的 reshape 同样需要小心处理源码注释也特别提到 2×2 盒均值直接使用avg_pool2d是因为在两个 unfold 维度上求均值虽然数值一致但在 MPS 上慢数倍——可见这个函数的实现本身就充分考虑了 MPS 行为。四、修复后的空输入语义空进空出empty-in → empty-out修复后Kornia 对空输入的契约是明确的空进空出。以逆向转换为例测试 tests/color/test_yuv.pyYuv420ToRgb与 tests/color/test_yuv.pyYuv422ToRgb分别钉死了这一约定def test_empty_input(self, device, dtype): # Regression for #4056: a *consistently* zero-sized luma/chroma pair passes the guard # (0 2 * 0) and returns an empty RGB plane, matching the 4:4:4 twin, rather than # raising. Pin the empty-in - empty-out convention instead of special-casing it away. y torch.rand(1, 0, 0, devicedevice, dtypedtype) uv torch.rand(2, 0, 0, devicedevice, dtypedtype) out kornia.color.yuv420_to_rgb(y, uv) assert out.shape (3, 0, 0) # agrees with the 4:4:4 converter on the same empty concatenation assert torch.equal(kornia.color.yuv_to_rgb(torch.cat([y, uv], dim-3)), out)这里传达了几个关键点一致性空对consistently zero-sized pair能通过守卫形状守卫判断0 2 * 0成立即亮度 0×0 色度 0×0被视为合法输入返回空 RGB 平面(3, 0, 0)而不是抛异常与 4:4:4 孪生转换yuv_to_rgb保持一致对同样的空拼接torch.cat([y, uv], dim-3)yuv_to_rgb的输出与yuv420_to_rgb完全相等torch.equal说明空语义在整族 YUV 转换中是统一的用测试钉死约定注释强调 Pin the empty-in → empty-out convention instead of special-casing it away即这是有意的 API 契约而不是实现附带的偶然行为。test_empty_input使用pytest.mark.parametrize在device与dtype两个维度上展开意味着 CPU/MPS/CUDA 各设备、各 dtype 都要满足该约定——这正是 #4185 修复所保证的设备一致性。五、不能忽视的兄弟修复空维度必须是 ShapeError#4056与 #4185 紧密相关的是更早的 #4056 回归修复在同一批测试注释中被反复提及。区别在于一致性空对luma 与 chroma 同时为 0 尺寸→ 合法返回空输出不一致空对只有色度的某一维为 0如imguv torch.ones(2, 2, 0)→ 依然是形状违规必须抛ShapeError而不能抛裸的ZeroDivisionError。相关测试见 tests/color/test_yuv.py# Regression for #4056: a zero-sized chroma dimension is still a shape violation and must # be reported as ShapeError, not the bare ZeroDivisionError the old ratio guard threw. with pytest.raises(ShapeError): imgy torch.ones(1, 4, 4, devicedevice, dtypedtype) imguv torch.ones(2, 2, 0, devicedevice, dtypedtype) kornia.color.yuv420_to_rgb(imgy, imguv)这说明旧的宽高比守卫ratio guard在遇到除数为 0 时曾泄漏出ZeroDivisionError修复后统一改为显式的ShapeError。两条修复共同构成了完整的空输入语义输入情况行为依据空且一致的 luma/chroma 对如(1,0,0)(2,0,0)返回空 RGB(3,0,0)与 CPU/4:4:4 一致test_yuv.py某维为 0 但比例不一致如 chroma 宽为 0抛ShapeError不抛ZeroDivisionErrortest_yuv.py奇数宽/高、色度非半尺寸抛ShapeErroryuv.py六、MPS 严格失败清单中的位置从 200 条降到 14 条再到 12 条Kornia 对 MPS 的支持采用严格失败清单 xfail 追踪的工程化方式清单文件是 testing/known_failure_xfails/mps_float32.txt。该文件头部注释直接印证了 #4185 的来龙去脉torch 2.14.0 added MPS kernels for the eigen/QR/LU/SVD family and for both grid_sample backward passes, which took the baseline from 200 entries to 14.#4185 then fixed the two YUV420/YUV422 empty-input failures, leaving 12 entries.One cause remains: NotImplementedError aten::_linalg_eigvals has no MPS kernel (the 5-point essential solver).也就是说在 torch 2.14.0 之前MPS 上的已知失败清单有约 200 条torch 2.14.0 为特征值/QR/LU/SVD 族及grid_sample两个反向 pass 新增了 MPS kernel清单骤降到 14 条#4185 又修复了 YUV420/YUV422 的两个空输入失败清单进一步降到 12 条剩余的 12 条全部是tests/geometry/epipolar/test_essential.py中依赖aten::_linalg_eigvalsMPS kernel 的用例。同文件还记录了 #4185 之后的一次全量回放结果the post-#4185 full-suite replay completed with --xfail-known-failures:8327 passed, 4126 skipped, 95 deselected, 43 xfailed, 9 xpassedin 425.15s.这条基线数据可以帮助你理解 MPS 支持的整体健康度绝大多数测试在 MPS 上直接通过仅剩 12 个真正受限于 PyTorch MPS kernel 缺失的已知失败。CI 侧TESTING.md 说明 PR CI 会跑一个阻塞性的 MPS 测试腿tests-mps使用 GitHub 托管的macos-15Apple Silicon 镜像、最新的 PyTorch、float32精度。这意味着任何让 MPS 清单倒退的改动都会被 CI 拦截——#4185 的修复正是在这套机制下被验证并纳入了严格基线。七、完整的复现与验证路径7.1 复现修复前的错误行为如果你手上有一个修复前的 Kornia 版本在 Apple Silicon 的 MPS 设备上执行import torch import kornia y torch.rand(1, 0, 0, devicemps) uv torch.rand(2, 0, 0, devicemps) # 修复前加速器分支 reshape(-1, ...) 无法推断空张量的 batch 维度报错 out kornia.color.yuv420_to_rgb(y, uv)修复前空张量在 MPS 上会落入 conv2d 分支reshape(-1, 3, 0, 0)因无法从 0 个元素中推断 batch 而失败。7.2 修复后的行为当前仓库在当前仓库版本下同一段代码在 MPS 上返回空 RGB 平面与 CPU 行为完全一致import torch import kornia for device in (cpu, mps): # 仅当机器是 Apple Silicon 时 mps 才可用 y torch.rand(1, 0, 0, devicedevice) uv torch.rand(2, 0, 0, devicedevice) out420 kornia.color.yuv420_to_rgb(y, uv) out422 kornia.color.yuv422_to_rgb(y, uv) print(device, out420.shape, out422.shape) # 两者均为 torch.Size([3, 0, 0])非空输入的常规用法不受影响例如import torch import kornia img torch.rand(2, 3, 4, 6) # 宽高必须为偶数 y, uv kornia.color.rgb_to_yuv420(img) print(y.shape) # torch.Size([2, 1, 4, 6]) print(uv.shape) # torch.Size([2, 2, 2, 3]) y, uv kornia.color.rgb_to_yuv422(img) print(y.shape) # torch.Size([2, 1, 4, 6]) print(uv.shape) # torch.Size([2, 2, 4, 3]) recovered kornia.color.yuv420_to_rgb(y, uv) print(recovered.shape) # torch.Size([2, 3, 4, 6])也支持nn.Module风格调用RgbToYuv420、RgbToYuv422、Yuv420ToRgb、Yuv422ToRgb等不过要注意YUV420/YUV422 系列模块在源码中标注了ONNX_EXPORTABLE False见 yuv.py 与 yuv.py而 4:4:4 的RgbToYuv/YuvToRgb则声明了ONNX_DEFAULT_INPUTSHAPE [-1, 3, -1, -1]可以导出 ONNX。7.3 运行仓库自带测试仓库测试已覆盖本次修复的全部关键路径可在仓库根目录运行# 空输入语义empty-in - empty-out、ShapeError 守卫、往返一致性、gradcheck、jit、dynamo pytest tests/color/test_yuv.py -k empty_input or exception or forth_and_back or gradcheck or jit or dynamo具体来说test_empty_inputtest_yuv.py 与 test_yuv.py钉死空进空出约定并断言与 4:4:4 孪生转换结果逐元素相等test_exception如 test_yuv.py覆盖ShapeError分支含 #4056 回归用例test_forth_and_backtest_yuv.py验证 RGB ⇄ YUV420 往返误差在容差内test_unit_upsampling用整数除法索引构造期望升采样结果防止repeat_interleave用错轴或倍数test_gradcheck/test_jit/test_dynamo保证修复没有破坏可微性、TorchScript 与torch.compile兼容性。八、小结changelog.d/migration-103.fixed.md记录的这次修复表面上是MPS 上两个空输入用例背后其实是 Kornia 对跨后端行为一致性的严谨工程态度根因_apply_linear_transformation的 conv2d 性能分支依赖reshape(-1, ...)而空张量的 batch 维度无法从 0 个元素中唯一推断导致 MPS 报错修法把numel() 0并入 einsum 分支空张量在所有后端统一走天然支持空形状的路径从根上规避 reshape 歧义契约空且一致的 luma/chroma 对返回空 RGBempty-in → empty-out与 4:4:4 转换一致不一致的空维度则抛ShapeError验证测试钉死约定、MPS 严格失败清单记录基线200 → 14 → 12 条CI 的tests-mps腿保证清单不倒退。如果你正在为其他颜色转换函数或自定义算子做 MPS/CUDA 适配本案例是一个值得借鉴的模板不要假设空张量在所有 reshape 下都安全优先让空输入走不受形状推断影响的路径并用空进空出 异常类型双层测试把语义钉死。赞分享计算机视觉深度学习人工智能图像处理【免费下载链接】kornia 空间人工智能的几何计算机视觉库项目地址https://gitcode.com/kornia/kornia点击查看免费下载相关推荐Kornia 空目标几何变换深度解析修复 warp_affine、warp_perspective 与 remap 在 MPS 及零尺寸 dsize 下的崩溃Kornia 空目标几何变换深度解析修复 warp_affine 、 warp_perspective 与 remap 在 MPS 及零尺寸 dsize 下的计算机视觉深度学习人工智能图像处理Kornia RandomCrop3D 的跳过语义与输入保真changelog 4667 修复背后的实现解析Kornia RandomCrop3D 的跳过语义与输入保真changelog 4667 修复背后的实现解析 本文围绕 Kornia 仓库 changelog计算机视觉人工智能深度学习图像处理Kornia RandomMosaic 修复深度解析非方形输入下的 (H, W) 保序与 start_ratio_range 轴语义修正Kornia RandomMosaic 修复深度解析非方形输入下的 H, W 保序与 start_ratio_range 轴语义修正 本篇技术指南以 Korn计算机视觉深度学习人工智能图像处理创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表