丁提取的 CPU 半精度采樣修復(fù):float32 回退與批處理 grid_sample 優(yōu)化解析)
計(jì)算機(jī)視覺(jué)深度學(xué)習(xí)人工智能圖像處理【免費(fèi)下載鏈接】kornia 空間人工智能的幾何計(jì)算機(jī)視覺(jué)庫(kù)項(xiàng)目地址https://gitcode.com/kornia/kornia點(diǎn)擊查看免費(fèi)下載導(dǎo)讀本文圍繞 changelog.d/migration-110.fixed.md 記錄的一次關(guān)鍵修復(fù)展開(kāi)在 PyTorch ≤ 2.9 中float16/bfloat16的 CPUgrid_sample內(nèi)核在旋轉(zhuǎn)補(bǔ)丁跨越圖像邊界時(shí)存在越界讀取缺陷會(huì)返回 NaN 或遠(yuǎn)超圖像取值范圍的垃圾值。Kornia 的extract_patches_simple與extract_patches_from_pyramid通過(guò)「在所有設(shè)備上以 float32 采樣再轉(zhuǎn)回半精度」的方式繞開(kāi)該缺陷同時(shí)把逐圖像 Python 循環(huán)折疊為一次批處理的grid_sample并在torch.compile(fullgraphTrue)下消除了隨 batch 增長(zhǎng)的計(jì)算圖膨脹。讀完本文你將掌握這兩個(gè) LAF 補(bǔ)丁提取器的精度策略、內(nèi)存分塊機(jī)制以及它們與 torch.compile 的交互原理。背景兩個(gè) LAF 補(bǔ)丁提取函數(shù)Kornia 的局部特征LAFLocal Affine Frames管線中extract_patches_simple與extract_patches_from_pyramid是核心的采樣算子均定義于 kornia/feature/laf.py導(dǎo)出入口見(jiàn) kornia/feature/init.pyextract_patches_simple(img, laf, PS32, normalize_lafs_before_extractionTrue)實(shí)現(xiàn)直接按 LAF 定義的仿射框從原圖采樣不做平滑因此有較強(qiáng)混疊aliasing適合快速提取。extract_patches_from_pyramid實(shí)現(xiàn)從圖像金字塔的合適層級(jí)采樣尺度小于PS的 LAF 會(huì)被路由到仍能提供完整補(bǔ)丁的最粗層級(jí)避免混疊。兩個(gè)函數(shù)的輸入輸出約定一致圖像為(B, CH, H, W)LAF 為(B, N, 2, 3)輸出補(bǔ)丁為(B, N, CH, PS, PS)。它們都是純函數(shù)式 API供LAFOrienter、LAFAffNetShapeEstimator、LAFDescriptor等下游模塊復(fù)用相關(guān)修復(fù)鏈條見(jiàn) changelog.d/migration-111.fixed.md。問(wèn)題根因torch ≤ 2.9 的半精度 CPU 越界讀取本次修復(fù)的核心是一個(gè)上游PyTorch缺陷當(dāng)旋轉(zhuǎn)后的補(bǔ)丁跨越圖像邊界時(shí)torch ≤ 2.9 中float16/bfloat16的 CPUgrid_sample內(nèi)核會(huì)越界讀取out-of-bounds read從而返回 NaN 或遠(yuǎn)超出圖像取值范圍的垃圾值例如數(shù)量級(jí)1e4的隨機(jī)堆數(shù)據(jù)。原因在于旋轉(zhuǎn)補(bǔ)丁的采樣網(wǎng)格有相當(dāng)一部分落在圖像外這些坐標(biāo)由邊界填充padding_modeborder來(lái)補(bǔ)值半精度 CPU 內(nèi)核在這些邊界坐標(biāo)上沒(méi)有正確處理讀取了非法內(nèi)存。該問(wèn)題在測(cè)試中專門(mén)以「補(bǔ)丁必須保持在圖像取值范圍內(nèi)」這一不變式來(lái)刻畫(huà)。見(jiàn) tests/feature/test_laf.py 的test_border_patches_stay_in_range任何跨界補(bǔ)丁采到的值都應(yīng)是圖像像素值的凸組合因此必須落在[img.min(), img.max()]之內(nèi)且有限isfinite。而缺陷內(nèi)核會(huì)返回零、1e4量級(jí)數(shù)值或 NaN取決于堆內(nèi)存中恰好殘留的內(nèi)容。修復(fù)方案一全設(shè)備 float32 采樣回退修復(fù)的核心策略是把半精度輸入在采樣前統(tǒng)一提升到 float32采樣完成后再把補(bǔ)丁轉(zhuǎn)回原始精度。從源碼看這一邏輯體現(xiàn)在兩條關(guān)鍵鏈路上網(wǎng)格精度提升_grid_dtype將float16/bfloat16一律映射為torch.float32_promoted_grid_dtype進(jìn)一步對(duì)圖像與 LAF 做torch.promote_types后應(yīng)用該策略保證坐標(biāo)運(yùn)算不丟失任何一側(cè)的精度。圖像一次上采樣在 extract_patches_simple 的實(shí)現(xiàn) 中sample_img img.to(grid_dtype)在分塊循環(huán)之前完成避免每個(gè)塊重復(fù)對(duì)整幅圖做類型轉(zhuǎn)換_grid_sample_patches實(shí)現(xiàn)保證img與grid共享 dtype直接調(diào)用F.grid_sample。值得注意的是該上采樣并非只針對(duì) torch ≤ 2.9 的壞內(nèi)核——_grid_sample_patches的 docstring 明確指出這是在所有 torch 版本上有意為之半精度采樣坐標(biāo)在大圖像上會(huì)量化為整像素即歸一化坐標(biāo)精度損失float32 網(wǎng)格計(jì)算可避免這一損失。_grid_dtype的 docstring 也說(shuō)明了它對(duì) float16/bfloat16 統(tǒng)一提升的原則。從代碼注釋看extract_patches_from_pyramid同樣遵循該約定金字塔 atlas 直接用網(wǎng)格 dtype 構(gòu)建kornia/feature/laf.py#L795-L807使 replicate 填充、pyrdown與每個(gè)分塊的grid_sample都運(yùn)行在所有 torch 版本都穩(wěn)定的內(nèi)核上。精度與性能的權(quán)衡源碼明確記錄了在 CUDA 上的代價(jià)原生半精度核在 CUDA 上是正常的不越界但為了統(tǒng)一的正確性策略float16 的extract_patches_simple在高 N 場(chǎng)景下大約付出2 倍減速換取準(zhǔn)確性。這在 migration-110 的 breaking changes 說(shuō)明 中被稱為 CUDA cost——即該修復(fù)改變了所有后端的半精度補(bǔ)丁數(shù)值而不僅是 CPU?;旌暇?autocast 管線的兼容測(cè)試test_mixed_dtype_laf_under_autocasttests/feature/test_laf.py#L753-L766驗(yàn)證了典型的 autocast 場(chǎng)景檢測(cè)器如 KeyNet在 autocast 下輸出 half/bfloat16 的 LAF而源圖仍是 float32。提取器自己完成小張量LAF的類型提升而不是拒絕這一合法管線或依賴后端各自的grid_sample隱式提升。對(duì)應(yīng)的test_mixed_dtype_preserves_laf_precisiontests/feature/test_laf.py#L768-L778則反向驗(yàn)證float32 LAF 配 half 圖像時(shí)LAF 的亞像素坐標(biāo)不能被降成 half 再提升回來(lái)——這正是網(wǎng)格使用提升后 dtype 的原因。修復(fù)方案二折疊的批處理 grid_sample此前兩個(gè)提取器采用「逐圖像 Python 循環(huán)」每張圖像調(diào)用一次grid_sample。本次修復(fù)將逐圖像循環(huán)替換為折疊的批處理grid_sample每個(gè) LAF 的(PS, PS, 2)網(wǎng)格被折疊為(B, N*PS, PS, 2)一次調(diào)用覆蓋整個(gè) batch。該邏輯見(jiàn)_sample_patchesfolded grid.view(B, N * PS, PS, 2)后單次_grid_sample_patches結(jié)果再 reshape 為(B, ch, N, PS, PS)并 permute 成(B, N, ch, PS, PS)。當(dāng)N很大時(shí)網(wǎng)格與采樣結(jié)果按N方向分塊chunking以約束工作區(qū)workspace峰值內(nèi)存。分塊大小由_grid_chunk_lafs決定它按B * PS * PS * max(2, ch) * elem_size估算每個(gè) LAF 的網(wǎng)格與通道縮放采樣結(jié)果的字節(jié)數(shù)用默認(rèn) 64 MiB 預(yù)算求最大可容納的 LAF 數(shù)當(dāng)一切都裝得下時(shí)循環(huán)退化為單次調(diào)用的快路徑。測(cè)試對(duì)該語(yǔ)義做了三重驗(yàn)證test_chunked_matches_single_calltests/feature/test_laf.py#L857-L870通過(guò) monkeypatch 把_grid_chunk_lafs強(qiáng)制為 1每塊一個(gè) LAF斷言與單次調(diào)用結(jié)果完全一致CPU 上逐位相等test_chunk_budget_accounts_for_channelstests/feature/test_laf.py#L872-L878驗(yàn)證高通道特征圖會(huì)把塊數(shù)壓小ch1時(shí) 1000 個(gè) LAF 一次裝下ch256時(shí)每塊僅 64 個(gè)test_channel_laf_correspondencetests/feature/test_laf.py#L880-L892逐 LAF 對(duì)比單塊提取結(jié)果鎖定ch1且N1時(shí)通道與 LAF 的對(duì)應(yīng)關(guān)系。對(duì) torch.compile(fullgraphTrue) 的圖優(yōu)化影響這是本次修復(fù)中容易被忽略卻影響深遠(yuǎn)的收益。遷移說(shuō)明中給出了具體的實(shí)證舊實(shí)現(xiàn)逐圖像循環(huán)在 trace 時(shí)會(huì)把循環(huán)展開(kāi)成每個(gè) batch 元素一個(gè)grid_sample計(jì)算圖隨 batch 增長(zhǎng)且每個(gè)新 batch 尺寸都會(huì)觸發(fā)一次重編譯。例如分別用 batch 2、3、5、7 追蹤extract_patches_simple會(huì)構(gòu)建出含 2/3/5/7 個(gè)grid_sample節(jié)點(diǎn)的 4 張圖新實(shí)現(xiàn)只需構(gòu)建各含 1 個(gè)節(jié)點(diǎn)的 2 張圖。換句話說(shuō)舊圖圖的規(guī)模 O(batch)batch 變化即重編譯新圖批處理折疊后圖固定為 1 個(gè)grid_sample節(jié)點(diǎn)圖的規(guī)模與 batch 解耦。兩種形式都能在fullgraphTrue下編譯但只有新實(shí)現(xiàn)避免了「圖隨 batch 線性膨脹 每次新 batch 重編譯」的開(kāi)銷。測(cè)試test_dynamotests/feature/test_laf.py#L926-L935直接以torch.compile(..., fullgraphTrue)斷言提取器可以完整 trace 為單張圖。無(wú)數(shù)據(jù)依賴分支的實(shí)現(xiàn)細(xì)節(jié)為了保證fullgraph路徑可編譯兩個(gè)提取器在非有限 LAF 處理上也刻意保持「無(wú)條件」extract_patches_simple在循環(huán)結(jié)束后統(tǒng)一masked_fill_清零kornia/feature/laf.py#L713-L715extract_patches_from_pyramid則對(duì)pyr_idx 0的條目做同樣的無(wú)條件清零kornia/feature/laf.py#L856-L872避免數(shù)據(jù)相關(guān)的 Python 分支導(dǎo)致圖斷裂。附帶修復(fù)非有限 LAF 與 MPS 邊界行為本次 migration 涉及的同族修復(fù)還包括migration-111任何含 NaN/Inf 的 LAF 幀哪怕只有中心一個(gè)元素在網(wǎng)格算術(shù)前被整體標(biāo)記并凈化返回全零補(bǔ)丁與零 LAF 梯度而不是把非法網(wǎng)格交給grid_sample——其 CPUgrid_sampler_2d_backward邊界填充內(nèi)核可能直接終止進(jìn)程。測(cè)試見(jiàn) tests/feature/test_laf.py#L780-L810。migration-126MPS 沒(méi)有padding_modeborder用「零填充 網(wǎng)格截?cái)唷鼓M時(shí)截?cái)嗄繕?biāo)從±1align_cornersFalse下是邊界像素的外邊緣會(huì)與零填充混出約一半的暗色值修正為最外層像素中心±(1 - 1/size)使跨界補(bǔ)丁與 CPU 的最大偏差從0.395降到2.5e-6。實(shí)際使用建議無(wú)需用戶側(cè)改動(dòng)float32回退對(duì)調(diào)用方透明輸入 half 圖像仍返回 half 補(bǔ)丁dtype 不變見(jiàn)test_border_patches_stay_in_range中對(duì)patches.dtype half_dtype的斷言。設(shè)備與 dtype 約定輸出補(bǔ)丁始終落在圖像所在設(shè)備與 dtype 上LAF 允許與圖像不同設(shè)備/精度提取器會(huì)先遷移kornia/feature/laf.py#L684-L688。測(cè)試test_laf_on_another_devicetests/feature/test_laf.py#L914-L924驗(yàn)證了該契約。大 batch 與大圖分塊機(jī)制默認(rèn)以 64 MiB 預(yù)算約束峰值內(nèi)存對(duì)超高通道特征圖如ch256建議預(yù)判塊數(shù)可用_grid_chunk_lafs估算避免意外的工作區(qū)壓力。訓(xùn)練管線若你的檢測(cè)器在訓(xùn)練中可能產(chǎn)出退化的非有限 LAF本次修復(fù)已保證此類幀得到零補(bǔ)丁與安全反向傳播無(wú)需額外防御邏輯。編譯部署若使用torch.compile(fullgraphTrue)批處理折疊后的提取器圖規(guī)模與 batch 解耦重編譯次數(shù)大幅下降可放心放入編譯后的推理/訓(xùn)練圖。總結(jié)extract_patches_simple與extract_patches_from_pyramid的這次修復(fù)同時(shí)解決了三個(gè)層面的問(wèn)題正確性繞開(kāi) torch ≤ 2.9 半精度 CPU 越界讀取、精度避免大圖像上歸一化坐標(biāo)的精度損失、性能/可編譯性批處理折疊 分塊內(nèi)存約束 消除隨 batch 膨脹的計(jì)算圖。其實(shí)現(xiàn)細(xì)節(jié)——提升后 dtype 的一次性上采樣、無(wú)數(shù)據(jù)依賴分支的清零策略、64 MiB/128 MiB 的分塊預(yù)算——都可在 kornia/feature/laf.py 與其配套測(cè)試 tests/feature/test_laf.py 中逐一驗(yàn)證是理解 Kornia 局部特征管線精度與性能權(quán)衡的絕佳案例。贊分享計(jì)算機(jī)視覺(jué)深度學(xué)習(xí)人工智能圖像處理【免費(fèi)下載鏈接】kornia 空間人工智能的幾何計(jì)算機(jī)視覺(jué)庫(kù)項(xiàng)目地址https://gitcode.com/kornia/kornia點(diǎn)擊查看免費(fèi)下載相關(guān)推薦kornia 非有限 LAF 幀修復(fù)深度解析零補(bǔ)丁輸出與 grid_sample 段錯(cuò)誤防護(hù)kornia 非有限 LAF 幀修復(fù)深度解析零補(bǔ)丁輸出與 grid_sample 段錯(cuò)誤防護(hù) 導(dǎo)讀 本文基于 kornia 倉(cāng)庫(kù) changelog 中關(guān)于計(jì)算機(jī)視覺(jué)深度學(xué)習(xí)人工智能圖像處理Kornia 修復(fù) MPS 邊界補(bǔ)丁提取暗化問(wèn)題grid_sample 邊框填充的像素中心鉗制原理Kornia 修復(fù) MPS 邊界補(bǔ)丁提取暗化問(wèn)題 grid_sample 邊框填充的像素中心鉗制原理 本篇文章圍繞 Kornia changelog 條目 c計(jì)算機(jī)視覺(jué)人工智能深度學(xué)習(xí)圖像處理Kornia LAF 特征點(diǎn)提取引擎重構(gòu)解析atlas 金字塔采樣、float32 網(wǎng)格精度與分塊內(nèi)存控制migration-012Kornia LAF 特征點(diǎn)提取引擎重構(gòu)解析atlas 金字塔采樣、float32 網(wǎng)格精度與分塊內(nèi)存控制migration 012 本文對(duì)應(yīng)倉(cāng)庫(kù) cha計(jì)算機(jī)視覺(jué)深度學(xué)習(xí)人工智能圖像處理創(chuàng)作聲明:本文部分內(nèi)容由AI輔助生成(AIGC),僅供參考