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

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

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

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

Arm NSS 通過(guò)一套從數(shù)據(jù)中學(xué)習(xí)且經(jīng)過(guò)訓(xùn)練的神經(jīng)模型，而非靜態(tài)規(guī)則來(lái)克服這些限制。它可以跨條件和內(nèi)容類(lèi)型進(jìn)行泛化，更有效地適應(yīng)運(yùn)動(dòng)動(dòng)力和識(shí)別鋸齒模式。這些能力使其能夠更可靠地處理那些對(duì) AMD 超級(jí)分辨率銳畫(huà)技術(shù) 2 (FSR2) 和 Arm 精銳超級(jí)分辨率技術(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) 圖像配對(duì)。序列大約包含 100 幀，以幫助模型理解圖像內(nèi)容如何隨時(shí)間變化。

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

該網(wǎng)絡(luò)采用時(shí)空損失函數(shù)進(jìn)行訓(xùn)練，能對(duì)空域保真度和時(shí)域一致性方面的誤差進(jìn)行懲罰。空域保真度用于確保單幀圖像清晰、細(xì)節(jié)豐富且視覺(jué)準(zhǔn)確，有助于保留邊緣、紋理和精細(xì)結(jié)構(gòu)；而時(shí)域穩(wěn)定性則抑制了連續(xù)幀之間可能出現(xiàn)的閃爍、抖動(dòng)或其他形式的時(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 語(yǔ)言編寫(xiě)，以確保靈活性和性能，而量化感知訓(xùn)練則采用 ExecuTorch 完成。

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

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

Arm 評(píng)估了多種方法：

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

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

參數(shù)預(yù)測(cè)（所選方法）?：為每個(gè)像素輸出少量參數(shù)。這些參數(shù)驅(qū)動(dòng)后處理步驟（如濾波和樣本累積）。該方法對(duì)量化友好且?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)化升級(jí)的最終圖像；

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

白皮書(shū)：https://developer.arm.com/documentation/111019

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

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

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

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

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

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

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

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

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