一文了解 Arm 神經(jīng)超級(jí)采樣,深入探索架構(gòu)、訓(xùn)練和推理
本文將從訓(xùn)練、網(wǎng)絡(luò)架構(gòu)到后處理和推理等方面,深入探討 Arm 神經(jīng)超級(jí)采樣 (Arm Neural Super Sampling, Arm NSS) 的工作原理,希望為機(jī)器學(xué)習(xí) (ML) 工程師和移動(dòng)端圖形開發(fā)者來詳細(xì)解釋 Arm NSS 的運(yùn)行機(jī)制,及其如何在移動(dòng)端硬件上進(jìn)行部署。
用神經(jīng)超級(jí)采樣取代啟發(fā)式方法的原因
時(shí)域超級(jí)采樣 (TSS),也稱為時(shí)域抗鋸齒 (TAA),已經(jīng)成為過去十年中抗鋸齒技術(shù)的行業(yè)標(biāo)準(zhǔn)解決方案。它具備諸多優(yōu)勢(shì),包括能夠解決各類鋸齒問題,對(duì)于延遲渲染的計(jì)算效率高,并且可擴(kuò)展至圖像的優(yōu)化升級(jí)。然而,該解決方案也面臨一些挑戰(zhàn)。例如,TSS 中目前常用的人工設(shè)計(jì)的啟發(fā)式方法難以擴(kuò)展,并且需要在不同內(nèi)容中不斷調(diào)整,像鬼影、去遮擋偽影和時(shí)域不穩(wěn)定性等問題依然存在,并且在與圖像優(yōu)化升級(jí)技術(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)動(dòng)動(dòng)力和識(shí)別鋸齒模式。這些能力使其能夠更可靠地處理那些對(duì) AMD 超級(jí)分辨率銳畫技術(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)行懲罰??沼虮U娑扔糜诖_保單幀圖像清晰、細(xì)節(jié)豐富且視覺準(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 語言編寫,以確保靈活性和性能,而量化感知訓(xùn)練則采用 ExecuTorch 完成。
網(wǎng)絡(luò)架構(gòu)和輸出設(shè)計(jì)
Arm NSS 網(wǎng)絡(luò)采用具有跳躍連接的四層級(jí) UNet 主干架構(gòu),以保持空間結(jié)構(gòu)。它通過三個(gè)編碼器和解碼器模塊對(duì)輸入數(shù)據(jù)進(jìn)行下采樣和上采樣處理。
Arm 評(píng)估了多種方法:
· 圖像預(yù)測(cè)?:易于實(shí)現(xiàn),但在量化條件下表現(xiàn)不佳,并會(huì)產(chǎn)生視覺偽影。
· 內(nèi)核預(yù)測(cè):泛化能力良好且能有效量化,但由于需要許多大型內(nèi)核映射,導(dǎo)致帶寬開銷較高。
· ?參數(shù)預(yù)測(cè)(所選方法)?:為每個(gè)像素輸出少量參數(shù)。這些參數(shù)驅(qū)動(dòng)后處理步驟(如濾波和樣本累積)。該方法對(duì)量化友好且?guī)捫矢摺?
該網(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 此類技術(shù)依賴手動(dòng)調(diào)優(yōu)的啟發(fā)式方法不同,Arm NSS 等 ML 方法具有三重優(yōu)勢(shì):
1. Arm NSS 能估計(jì)動(dòng)態(tài)內(nèi)核濾波器和參數(shù),以逐像素粒度解決鋸齒問題;
2. Arm NSS 通過利用采集多幀歷史狀態(tài)的時(shí)域反饋,實(shí)現(xiàn)更出色的時(shí)域穩(wěn)定性;
3. Arm NSS 可針對(duì)新游戲內(nèi)容進(jìn)行調(diào)優(yōu),使開發(fā)者能進(jìn)一步針對(duì)其特定游戲優(yōu)化圖像質(zhì)量。
通過時(shí)域反饋提升幀間一致性
Arm NSS 引入了兩種關(guān)鍵反饋機(jī)制來解決時(shí)域不穩(wěn)定性問題:
· 將先前幀的隱藏特征前向傳遞,使網(wǎng)絡(luò)能夠?qū)W習(xí)變化與持續(xù)的部分;
· 通過計(jì)算亮度導(dǎo)數(shù)來檢測(cè)閃爍的細(xì)微特征,從而突顯指示不穩(wěn)定的時(shí)域差異。
這些輸入幫助模型在不依賴人工設(shè)計(jì)規(guī)則的情況下保持時(shí)域穩(wěn)定性。
預(yù)處理階段:輸入準(zhǔn)備
在推理開始之前,基于 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)建最終輸出顏色。所有步驟均集成至渲染圖中,并專為移動(dòng)端高效運(yùn)行設(shè)計(jì),具體包括以下步驟:
· 運(yùn)動(dòng)向量擴(kuò)張,減少重投影歷史時(shí)的鋸齒問題;
· ?歷史重投影?,采用 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í)行插值和抗鋸齒處理,類似于去馬賽克;
· 校正,使用預(yù)測(cè)的 theta 參數(shù)剔除陳舊歷史數(shù)據(jù);
· 樣本累積?,通過預(yù)測(cè)的 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)來評(píng)估 Arm NSS。雖然這些指標(biāo)未必總是匹配人類感知,但它們有助于發(fā)現(xiàn)問題案例,并通過跟蹤各種指標(biāo)增強(qiáng)信心。
通過持續(xù)集成 (CI) 工作流回放測(cè)試序列并記錄 Arm NSS、Arm ASR 和其他基準(zhǔn)的性能。有關(guān)視覺比較和感知評(píng)估的詳細(xì)內(nèi)容,請(qǐng)參閱技術(shù)白皮書。
在 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 無需反應(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ù)的性能條件下,我們對(duì)每幀的優(yōu)化升級(jí)處理設(shè)定了?不超過四毫秒?的目標(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)化的推理過程取代了復(fù)雜的啟發(fā)式方法。
即刻利用 Arm NSS 進(jìn)行構(gòu)建
Arm NSS 引入了一種由 ML 驅(qū)動(dòng)的時(shí)域超級(jí)采樣的實(shí)用方法。它用經(jīng)過學(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 工程師來說,Arm NSS 是一個(gè)可部署、結(jié)構(gòu)良好,且運(yùn)行于圖形管線內(nèi)的推理示例。
想要探索 Arm 神經(jīng)圖形開發(fā)套件,查看示例代碼,并探索網(wǎng)絡(luò)結(jié)構(gòu),請(qǐng)?jiān)L問 Arm 開發(fā)者中心的 Arm NSS 頁面。歡迎開發(fā)者使用開發(fā)套件或?yàn)樽约旱膬?nèi)容重新訓(xùn)練 NSS,并提供反饋,你的見解將有助于塑造移動(dòng)設(shè)備上神經(jīng)渲染的未來。