作者：Arm 工程部計(jì)算機(jī)視覺高級架構(gòu)師 Liam O’Neil

本文將從訓(xùn)練、網(wǎng)絡(luò)架構(gòu)到后處理和推理等方面，深入探討 Arm 神經(jīng)超級采樣 (Arm Neural Super Sampling, Arm NSS) 的工作原理，希望為機(jī)器學(xué)習(xí) (ML) 工程師和移動端圖形開發(fā)者來詳細(xì)解釋 Arm NSS 的運(yùn)行機(jī)制，及其如何在移動端硬件上進(jìn)行部署。

用神經(jīng)超級采樣取代啟發(fā)式方法的原因

時(shí)域超級采樣 (TSS)，也稱為時(shí)域抗鋸齒 (TAA)，已經(jīng)成為過去十年中抗鋸齒技術(shù)的行業(yè)標(biāo)準(zhǔn)解決方案。它具備諸多優(yōu)勢，包括能夠解決各類鋸齒問題，對于延遲渲染的計(jì)算效率高，并且可擴(kuò)展至圖像的優(yōu)化升級。然而，該解決方案也面臨一些挑戰(zhàn)。例如，TSS 中目前常用的人工設(shè)計(jì)的啟發(fā)式方法難以擴(kuò)展，并且需要在不同內(nèi)容中不斷調(diào)整，像鬼影、去遮擋偽影和時(shí)域不穩(wěn)定性等問題依然存在，并且在與圖像優(yōu)化升級技術(shù)結(jié)合時(shí)變得更加棘手。

Arm NSS 通過一套從數(shù)據(jù)中學(xué)習(xí)且經(jīng)過訓(xùn)練的神經(jīng)模型，而非靜態(tài)規(guī)則來克服這些限制。它可以跨條件和內(nèi)容類型進(jìn)行泛化，更有效地適應(yīng)運(yùn)動動力和識別鋸齒模式。這些能力使其能夠更可靠地處理那些對 AMD 超級分辨率銳畫技術(shù) 2 (FSR2) 和 Arm 精銳超級分辨率技術(shù) (Arm Accuracy Super Resolution, Arm ASR) 等方法頗具挑戰(zhàn)的邊緣側(cè)用例。

訓(xùn)練 Arm NSS 網(wǎng)絡(luò)：帶反饋的遞歸學(xué)習(xí)

Arm NSS 使用以每像素一個(gè)樣本渲染的 540p 幀序列進(jìn)行訓(xùn)練。每幀與以每像素 16 個(gè)樣本渲染的 1080p 真實(shí)標(biāo)簽 (Ground Truth) 圖像配對。序列大約包含 100 幀，以幫助模型理解圖像內(nèi)容如何隨時(shí)間變化。

解決方案的輸入包括渲染圖像的顏色、運(yùn)動向量和深度，以及引擎元數(shù)據(jù)，如抖動向量和攝像頭矩陣。該模型采用遞歸式訓(xùn)練，在每次執(zhí)行反向傳播之前，會先對多幀序列進(jìn)行前向傳播。這使得網(wǎng)絡(luò)能夠隨時(shí)間傳播梯度，并學(xué)習(xí)如何累積信息。

該網(wǎng)絡(luò)采用時(shí)空損失函數(shù)進(jìn)行訓(xùn)練，能對空域保真度和時(shí)域一致性方面的誤差進(jìn)行懲罰?？沼虮Ｕ娑扔糜诖_保單幀圖像清晰、細(xì)節(jié)豐富且視覺準(zhǔn)確，有助于保留邊緣、紋理和精細(xì)結(jié)構(gòu)；而時(shí)域穩(wěn)定性則抑制了連續(xù)幀之間可能出現(xiàn)的閃爍、抖動或其他形式的時(shí)域噪聲。

此項(xiàng)訓(xùn)練是在 PyTorch 中采用成熟的技術(shù)實(shí)踐進(jìn)行，其中包括 Adam 優(yōu)化器、余弦退火學(xué)習(xí)率調(diào)度，以及標(biāo)準(zhǔn)數(shù)據(jù)增強(qiáng)策略。預(yù)處理和后處理流程使用 Slang 語言編寫，以確保靈活性和性能，而量化感知訓(xùn)練則采用 ExecuTorch 完成。

網(wǎng)絡(luò)架構(gòu)和輸出設(shè)計(jì)

Arm NSS 網(wǎng)絡(luò)采用具有跳躍連接的四層級 UNet 主干架構(gòu)，以保持空間結(jié)構(gòu)。它分別通過三個(gè)編碼器和解碼器模塊對輸入數(shù)據(jù)進(jìn)行下采樣和上采樣處理。

Arm 評估了多種方法：

圖像預(yù)測?：易于實(shí)現(xiàn)，但在量化條件下表現(xiàn)不佳，并會產(chǎn)生視覺偽影。

內(nèi)核預(yù)測：泛化能力良好且能有效量化，但由于需要許多大型內(nèi)核映射，導(dǎo)致帶寬開銷較高。

參數(shù)預(yù)測（所選方法）?：為每個(gè)像素輸出少量參數(shù)。這些參數(shù)驅(qū)動后處理步驟（如濾波和樣本累積）。該方法對量化友好且?guī)捫矢摺?/p>

該網(wǎng)絡(luò)生成三種逐像素輸出：

4 x 4 濾波器內(nèi)核；

用于累積和校正的時(shí)域系數(shù)?；

隱藏狀態(tài)張量?作為時(shí)域反饋傳遞至下一幀。

該網(wǎng)絡(luò)輸出服務(wù)于兩條路徑：

濾波器內(nèi)核與時(shí)域系數(shù)?被后處理階段用于計(jì)算經(jīng)優(yōu)化升級的最終圖像；

隱藏狀態(tài)被前向傳遞至下一幀推理。與 Arm ASR 此類技術(shù)依賴手動調(diào)優(yōu)的啟發(fā)式方法不同，Arm NSS 等 ML 方法具有三重優(yōu)勢：

Arm NSS 能估計(jì)動態(tài)內(nèi)核濾波器和參數(shù)，以逐像素粒度解決鋸齒問題；

Arm NSS 通過利用采集多幀歷史狀態(tài)的時(shí)域反饋，實(shí)現(xiàn)更出色的時(shí)域穩(wěn)定性；

Arm NSS 可針對新游戲內(nèi)容進(jìn)行調(diào)優(yōu)，使開發(fā)者能進(jìn)一步針對其特定游戲優(yōu)化圖像質(zhì)量。

通過時(shí)域反饋提升幀間一致性

Arm NSS 引入了兩種關(guān)鍵反饋機(jī)制來解決時(shí)域不穩(wěn)定性問題：

將先前幀的隱藏特征前向傳遞，使網(wǎng)絡(luò)能夠?qū)W習(xí)變化與持續(xù)的部分；

通過計(jì)算亮度導(dǎo)數(shù)來檢測閃爍的細(xì)微特征，從而突顯指示不穩(wěn)定的時(shí)域差異。

這些輸入幫助模型在不依賴人工設(shè)計(jì)規(guī)則的情況下保持時(shí)域穩(wěn)定性。

預(yù)處理階段：輸入準(zhǔn)備

在推理開始之前，基于 GPU 的預(yù)處理階段會準(zhǔn)備 Arm NSS 所需的輸入數(shù)據(jù)，其中包括收集每個(gè)像素的屬性（如顏色、運(yùn)動向量和深度）。此外，該階段還會計(jì)算亮度導(dǎo)數(shù)，這是一個(gè)用于標(biāo)記細(xì)微特征閃爍的時(shí)域信號，以及去遮擋掩碼，以突出顯示陳舊的歷史記錄，并重投影歷史隱藏特征。

這些數(shù)據(jù)將被整合為用于神經(jīng)網(wǎng)絡(luò)的單一輸入張量。該階段作為計(jì)算著色器運(yùn)行，在推理調(diào)用前執(zhí)行，即在使用面向 Vulkan 的 ML 擴(kuò)展的 GPU 上運(yùn)行。

后處理：從原始輸出到最終幀

在推理完成后，作為計(jì)算著色器運(yùn)行的后處理階段將構(gòu)建最終輸出顏色。所有步驟均集成至渲染圖中，并專為移動端高效運(yùn)行設(shè)計(jì)，具體包括以下步驟：

運(yùn)動向量擴(kuò)張，減少重投影歷史時(shí)的鋸齒問題；

歷史重投影?，采用 Catmull-Rom 濾波器降低重投影的模糊度；

濾波，應(yīng)用 4 x 4 內(nèi)核對當(dāng)前顏色輸入執(zhí)行抗鋸齒處理；

稀疏優(yōu)化升級，將抖動的低分辨率樣本映射到高分辨率網(wǎng)格上，缺失像素都以零值填充，然后用 4 x 4 內(nèi)核稀疏濾波，執(zhí)行插值和抗鋸齒處理，類似于去馬賽克；

校正，使用預(yù)測的 theta 參數(shù)剔除陳舊歷史數(shù)據(jù)；

樣本累積?，通過預(yù)測的 alpha 參數(shù)將新數(shù)據(jù)與歷史緩沖混合，在色調(diào)映射域執(zhí)行以避免“螢火蟲”偽影。

驗(yàn)證質(zhì)量

Arm 使用了峰值信噪比 (Peak Signal-to-Noise Ratio, PSNR)、結(jié)構(gòu)相似性指數(shù) (Structural Similarity Index, SSIM) 和渲染聚焦感知誤差指標(biāo) FLIP 等指標(biāo)來評估 Arm NSS。雖然這些指標(biāo)未必總是匹配人類感知，但它們有助于發(fā)現(xiàn)問題案例，并通過跟蹤各種指標(biāo)增強(qiáng)信心。

通過持續(xù)集成 (CI) 工作流回放測試序列并記錄 Arm NSS、Arm ASR 和其他基準(zhǔn)的性能。有關(guān)視覺比較和感知評估的詳細(xì)內(nèi)容，請參閱技術(shù)白皮書。

白皮書：https://developer.arm.com/documentation/111019

在 540p 到 1080p 的比較中，Arm NSS 提升了穩(wěn)定性和細(xì)節(jié)保留性能，并且在快速運(yùn)動、部分遮擋物體和細(xì)微幾何體場景中表現(xiàn)出色。與 Arm ASR 或 AMD FSR2 等非神經(jīng)方法不同，Arm NSS 無需反應(yīng)掩碼，即可處理粒子效果。

Arm NSS 能否實(shí)時(shí)運(yùn)行？

雖然搭載神經(jīng)加速器的芯片尚未發(fā)布，但我們可以通過一些最低性能假設(shè)和執(zhí)行網(wǎng)絡(luò)推理所需的乘積累加運(yùn)算 (MAC) 數(shù)量來估計(jì) Arm NSS 是否足夠快。該分析適用于任何在吞吐量、功率和利用率方面滿足這些相同假設(shè)的加速器。我們假定在可持續(xù)的 GPU 時(shí)鐘頻率下，目標(biāo)每瓦 10 TOP/s 的神經(jīng)加速是可實(shí)現(xiàn)的。

在持續(xù)的性能條件下，我們對每幀的優(yōu)化升級處理設(shè)定了?不超過四毫秒?的目標(biāo)。在低時(shí)鐘頻率 GPU 上，推理前后運(yùn)行的著色器階段共耗時(shí)約 ?1.4 毫秒??；诖祟A(yù)算，Arm NSS 需控制在約 ?27 GOPs ?以下。而參數(shù)預(yù)測的網(wǎng)絡(luò)使用約為 ?10 GOPs?，即使神經(jīng)網(wǎng)絡(luò)加速器的效率僅為 40%，這在該范圍內(nèi)仍綽綽有余。

早期模擬數(shù)據(jù)顯示，Arm NSS 在 1.5 倍優(yōu)化升級（平衡模式）下的運(yùn)行時(shí)間約為 Arm ASR的 75%，而在兩倍優(yōu)化升級（平衡模式）下則預(yù)計(jì)性能會優(yōu)于 Arm ASR。這一效率提升得益于用簡化的推理過程取代了復(fù)雜的啟發(fā)式方法。

即刻利用 Arm NSS 進(jìn)行構(gòu)建

Arm NSS 引入了一種由 ML 驅(qū)動的時(shí)域超級采樣的實(shí)用方法。它用經(jīng)過學(xué)習(xí)的濾波器和穩(wěn)定性提示取代了人工調(diào)優(yōu)的啟發(fā)式方法，并能夠在移動端硬件的實(shí)時(shí)性限制內(nèi)運(yùn)行。

其訓(xùn)練方法、緊湊架構(gòu)和利用面向 Vulkan 的 ML 擴(kuò)展使其兼具高性能和適應(yīng)性。對于構(gòu)建神經(jīng)渲染解決方案的 ML 工程師來說，Arm NSS 是一個(gè)可部署、結(jié)構(gòu)良好，且運(yùn)行于圖形管線內(nèi)的推理示例。

想要探索 Arm 神經(jīng)圖形開發(fā)套件，查看示例代碼，并探索網(wǎng)絡(luò)結(jié)構(gòu)，請?jiān)L問 Arm 開發(fā)者中心的 Arm NSS 頁面。歡迎開發(fā)者使用開發(fā)套件或?yàn)樽约旱膬?nèi)容重新訓(xùn)練 NSS，并提供反饋，你的見解將有助于塑造移動設(shè)備上神經(jīng)渲染的未來。