背景：
近年來，深度學習研究和産業化得到迅猛發展。随着移動端算力提升，實時性需求，以及使用者資料隐私方面的考量，導緻越來越多的推理任務從雲端轉移到移動端。移動平台上深度學習推理需要解決硬體平台、驅動、編譯優化、模型壓縮、算子算法優化和部署等問題。許多企業正在開發面向移動端的開源深度學習架構，如小米MACE、騰訊優圖NCNN/TNN、阿裡MNN和谷歌TensorFlow Lite等。現有的推理平台通常存在與第三方庫和硬體平台驅動相關的适配、模型驗證，資料互動等問題，導緻其優化、部署、維護和更新繁瑣。
對于移動端的圖形圖像類AI應用，由于計算能力、低功耗、I/O帶寬、記憶體架構和通路方式等限制，現有的推理架構很難滿足産品對實時性、性能和功耗的嚴格要求。為滿足移動端圖形圖像後處理高效AI推理的需求，提高效能并降低成本，我們在《面向AI圖形的輕量深度學習推理引擎 ShaderNN》[1]一文介紹了基于OpenGL 後端的GPU Shader的高效推理引擎ShaderNN一期項目，在這篇文章中，我們将介紹ShaderNN 2.0提供的新增功能，主要技術特性（高性能，輕量性，通用性和易擴充）、性能功耗優勢以及一些典型的應用場景。
目前主流移動端處理器的性能提升，除了得益于CPU主頻提升和多核化以外，更多是通過配置多核GPU硬體作為圖形渲染和圖像處理加速器，例如ARM Mali GPU，Qualcomm Adreno GPU和Apple A15。主流移動端推理架構也基本都提供了基于GPU的推理加速，其中廣泛采用的GPU程式設計接口有OpenGL、Vulkan和OpenCL。移動端主流的圖形及圖像處理類應用主要是基于OpenGL 或 Vulkan 圖形棧開發的。推理架構如果采用其他加速方案，如NPU/DSP、OpenCL，在與上述應用進行整合時，面臨的主要問題就是資料互動方面的開銷。ShaderNN 一期推理架構基于OpenGL，可以直接處理OpenGL texture并輸出OpenGL texture。Vulkan作為新一代的圖形棧，是一個跨平台的底層圖形和計算API，而OpenGL是一套更進階的圖形API。下面是Vulkan相對于OpenGL的幾個優勢：
- 更低的開銷：Vulkan設計的目标是提供更低的CPU開銷和更高的多線程性能。它允許開發人員更好地控制硬體資源，如顯存和線程，進而減少驅動程式的開銷，提高應用程式的性能。
- 顯式控制：Vulkan要求開發人員更加顯式地管理圖形渲染管線和資源。這種顯式控制使開發人員能夠更好地優化和管理資源，以适應不同的硬體平台。
- 支援多線程：Vulkan支援多線程渲染，允許開發人員更好地利用多核處理器。這可以提高圖形和計算任務的性能，并實作更好的并行處理。
- 更好的記憶體管理：Vulkan提供了更靈活的記憶體管理機制，使開發人員能夠更好地控制記憶體配置設定和使用。這使得開發人員可以更有效地管理顯存，減少記憶體碎片化，并優化應用程式的性能。
- 更好的異步支援：Vulkan允許開發人員在不同的任務之間進行更好的異步處理。這使得開發人員可以更好地利用GPU的計算能力，同時執行多個任務，提高應用程式的性能和效率。
總的來說，Vulkan相對于OpenGL提供了更低的開銷、更好的多線程支援、更靈活的記憶體管理和更好的異步處理能力。這些特性使得Vulkan在需要更高性能和更底層控制的圖形圖像應用程式中成為更好的選擇。Vulkan還推出一些高性能特性如光線追蹤擴充等的原生支援，受到越來越多的軟體和硬體廠商的歡迎和支援，被應用于最新的遊戲、圖形圖像等應用。為滿足越來越多的移動端基于Vulkan應用的AI後處理的需求，它同樣需要一款輕量、可定制、可與應用Vulkan渲染管線無縫銜接的推理引擎。面向這一需求，我們在一期ShaderNN 1.0 的基礎上，添加了Vulkan 後端的支援，推出了首個支援OpenGL 和 Vulkan全圖形棧 Shader 2.0移動端推理引擎，并于近期釋出了開源的預覽版。
一、ShaderNN 2.0 流程圖
圖1: ShaderNN 工作流程圖
如圖1所示，ShaderNN 整個工作流程可以分為模型轉換和層融合優化、網絡模型和權重加載解析、計算圖生成、計算圖算子的執行、推理預測傳回相關結果。在編譯階段和運作階段進行優化，推理引擎分别進行了優化。編譯的優化包括Shader的編譯和緩存優化，等效算子替換與融合等。在運作階段，ShaderNN通過卷積優化、紋理重用、CPU和GPU 記憶體重用、資料結構布局優化、緩存及向量化優化以提升效能。
TensorFlow、或者PyTorch訓練後導出的ONNX格式的模型，通過ShaderNN轉換工具轉換成ShaderNN支援的 JSON模型格式。在轉換過程中，模型轉換器解耦模型結構和權重，解析算子并做層間融合優化。在加載模結構和權重以後，推理引擎通過拓撲排序生成計算圖。除少數算子通過CPU來實作以外，大多數算子交由GPU Shader來執行，在1.0中我們通過OpenGL後端的 Compute Shader 或者 Fragment Shader 來實作，在 2.0 中添加了Vulkan後端的 Compute Shader實作。
圖2: ShaderNN 算法時序圖
引擎的核心在MixedInferenceCore，它支援OpenGL和Vulkan後端，負責推理算子的具體實作。計算圖的輸入直接來自應用的圖形紋理（texture），它的輸出也是紋理，可供應用直接渲染。計算圖上的每一個節點對應網絡模型的一個Layer，每個Layer通過一個render stage 來計算。針對compute shader 和 fragment shader的管線不同，對于OpenGL fragment shader，每個render stage 包含一個或多個render pass，對于OpenGL 和 Vulkan compute shader，一個render stage 包含一個或多個render pass。Render pass最後的結果直接輸出到紋理（texture）。
二、ShaderNN的創新
- 使用基于紋理的輸入/輸出，提供與實時圖形管道或圖像處理應用程式的高效、零拷貝內建，進而節省 CPU 和 GPU 之間昂貴的資料傳輸和格式轉換，這對于移動平台上的實時應用程式至關重要。
圖3：ShaderNN同其它推理引擎同圖形圖像管線內建的對比
圖3對比了ShaderNN和其它推理引擎與現有管線內建的不同，其它推理引擎通常提供了稱為 Tensor 或 Mat 的通用資料結構作為輸入和輸出的格式，無法使用廣泛用于圖形和圖像應用程式的紋理作為輸入輸出，需要進行格式轉換或者GPU到CPU的來回轉換，對于圖像圖像類的大輸入的實時應用，這部分的開銷是相當大的。而ShaderNN以原生紋理作為輸入輸出沒有這部分開銷，大大的提升了整體的性能。
2.首個支援基于 OpenGL 後端Fragment Shader推理實作，對于輸入尺寸大、網絡層次簡單的網絡如超分、降噪網絡等相對其它推理引擎有顯著的優勢。
3.基于原生 OpenGL ES 和 Vulkan 建構，可輕松與圖形渲染管線內建，最大限度地利用計算資源，适用于對實時性要求高渲染、圖像/視訊和遊戲 AI 應用。
4.支援Compute Shader 和 Fragment Shader的混合實作，并能夠對不同的模型層選擇不同Shader以進行性能優化。友善使用者定制新的算子。
5.純GPU shader實作，無需依賴第三方庫，可在不同GPU硬體平台上使用，可定制，友善優化、內建、部署和更新。
三、 性能和功耗
在《面向AI圖形的輕量深度學習推理引擎 ShaderNN》[1]一文中，我們在4款手機同TensorFlow Lite OpenGL 後端對4個CNN常用網絡進行了單次推理時間和功耗上進行了對比。在推理時間上，在Spatial Denoise和ESPCN方面優于 TensorFlow Lite 75%-90%；Resnet18 和 YOLO V3 Tiny模型在部分處理器晶片上優于TensorFlow Lite 50%。在功耗方面，Spatial Denoise、ESPCN、Resnet18 和 YOLO V3 Tiny對比時，可節省高達 80%、70%、55% 和 51%。
對于ShaderNN 2.0 裡添加的Vulkan後端，我們和同樣支援Vulkan後端的MNN推理引擎,分别在兩款聯發科和兩款高通平台（見圖4）上, 挑選了Resnet18, Spatial Denoise，Style Transfer， ESPCN 4種CNN常用模型, 進行了單次推理時間和功耗上進行了對比。
圖4: 測試用移動平台
1、ShaderNN Vulkan 後端和 MNN Vulkan 後端性能對比
圖5: 性能對比
從圖5性能對比上看，在高通及MTK特定晶片組上表現優異，在Spatial Denoise和ESPCN等任務上的性能比MNN提高了50％-80％，在Resnet18和Style Transfer等任務上性能提高了6％-60％。
2、ShaderNN Vulkan 後端和 MNN Vulkan 後端功耗對比
圖6: 功耗對比
從圖6功耗對比上看，在執行Spatial Denoise、ESPCN、Resnet18和Style Transfer等推理任務時，ShaderNN Vulkan後端相對MNN可以分别節省高達60％、70％、45％和70％的能量消耗。
四、 ShaderNN的典型場景
從上面性能和功耗功耗對比可以看出，ShaderNN擅長圖形、圖像類的應用。ShaderNN 2.0針對圖形類應用場景，例如光線追蹤降噪、遊戲深度學習超采樣MNSS，進行了适配和管線優化。同時，通過添加新的算子，ShaderNN 也能支援AIGC類的應用，如 Stable Diffusion。
移動端實時風格遷移應用：
在《面向AI圖形的輕量深度學習推理引擎 ShaderNN》[1]一文中我們展示了風格遷移的Android Demo App，這個App 通過一個實時管線展示如何整合ShaderNN推理引擎到Android App，以及運作Style Transfer模型來處理從手機相機擷取的實時視訊流，并輸出風格遷移後的視訊流。在該應用中， 所有的資料都在GPU端處理，沒有額外的資料傳輸和處理消耗。
圖7: ShaderNN 風格遷移應用
移動端光線追蹤降噪應用：
光線追蹤（Path Tracing）是一種用于逼真渲染圖像的算法。它模拟光線從相機出發，經過場景中的物體，最終到達光源的路徑。通過追蹤每條光線的路徑并計算光線與物體的交點和光的傳播，可以生成高品質、逼真的圖像。和傳統的光栅化渲染相比，光線追蹤引入的計算量會成倍增加。光線追蹤采用一種随機蒙特卡洛采樣的方法，是以生成的圖像受到采樣噪聲的影響。離線渲染主要采用多次采樣的方式來降低噪聲，然而，在移動端實時渲染中，低采樣帶來的噪聲尤為明顯，該問題隻能通過降噪後處理來改善。
Intel推出了基于autoencoder結構的深度學習模型 Intel Open Image Denoise (OIDN)（https://github.com/OpenImageDenoise/oidn）。OIDN以反照率 （albedo）， 第一次反射的法線 （first bounce normals） 和噪聲資料 (noisy RGB image) 作為輸入，通過深度學習模型進行降噪，在PC端已取得良好的效果。ShaderNN 2.0添加了多輸入的支援，将這一模型應用到移動端圖形渲染管線中。圖8描述了降噪模型的基本流程，圖9中展示了降噪模型中實時輸入和輸出效果。
圖8：光線追蹤降噪子產品
輸入：1SPP HDR (帶噪聲)
輸入：Normal
輸出：降噪後RGB
輸入：Albedo
圖9：ShaderNN 光線追蹤降噪效果
移動端遊戲的深度學習超級采樣應用：
深度學習超采樣DLSS（Deep Learning Super Sampling）是由NVIDIA開發的一種圖像重建技術，旨在通過較低分辨率渲染結果和深度學習模型來建構高分辨率圖像，以達到提高圖像品質和性能的目的。通過ShaderNN這個開源項目，我們與浙江大學合作，開發實作了移動端實時深度學習超采樣模型[3]，并通過ShaderNN推理加速，首次在移動端實作遊戲實時超采樣。目前主流的DLSS模型都通過子采樣技術擷取更多像素細節，以及曆史幀和運動矢量矯正後的temporal時态資訊，來重建新的高分辨率圖像。目前幀建構推理的輸出結果又會回報給超采樣模型，作為下一幀建構的模型輸入。基于ShaderNN，移動端DLSS模型推理針對移動端進行了大量的性能優化，既滿足了實時性需求，又改善了遊戲畫面品質，尤其是在邊緣和細節部分。圖10描述了移動端DLSS模型的基本流程，圖11中可以看到模型的實際輸出效果。
圖10：移動端深度學習超采樣子產品
圖11：ShaderNN 運作移動端深度學習超采樣效果
移動端Stable Diffusion應用：
文字生成圖像領域最前沿的Stable Diffusion模型包含CLIP Text Encoder, UNet, VAE 三個部分，我們嘗試在ShaderNN2.0中運作Stable Diffusion模型https://huggingface.co/lambdalabs/miniSD-diffusers，并且部署到移動端。圖12展示了運作效果。
圖12：ShaderNN AIGC應用Stable Diffusion
五、規劃和展望
基于OpenGL後端的ShaderNN 1.0作為開源項目釋出後，ShaderNN 2.0添加了基于Vulkan後端的支援，支援移動端全圖形棧的AI應用，支援圖形圖像領域的最新應用，如風格遷移、Ray Tracing Denoising （光追降噪）、遊戲的深度學習超級采樣、Stable Diffusion等，展現ShaderNN在與圖形圖像類應用進行深度耦合方面的優勢。
我們下一步希望通過開源社群能吸引更多的開發者參與進來，覆寫更多的場景，持續優化算子和模型，以便更多的應用可以實際落地到産品中，緻力于共同打造一個具有特色的基于OpenGL、Vulkan的全圖形棧的移動推理引擎。
GitHub開源位址 （Apache2.0開源許可）
https://github.com/inferenceengine/shadernn
參考文獻：
1. 面向AI圖形的輕量深度學習推理引擎 ShaderNN，https://zhuanlan.zhihu.com/p/579051507
2.Justin Johnson, Alexandre Alahi, and Li Fei-Fei. 2016. Perceptual losses for real-time style transfer and super-resolution. In European conference on computer vision. Springer, 694–711
3. Sipeng Yang, Yunlu Zhao, Yuzhe Luo, He Wang, Hongyu Sun, Chen Li, Binghuang Cai, Xiaogang Jin,"MNSS: Neural Supersampling Framework for Real-Time Rendering on Mobile Devices", IEEE TRANSACTIONS ON VISUALIZATION AND COMPUTER GRAPHICS, Citation information: DOI 10.1109/TVCG.2023.3259141
本文作者：Qiang Qiu, Hongyu Sun, Yuzhong Yan, OPPO Computing & Graphics Research Institute
本文首發在“核心工匠”微信公衆号，歡迎關注公衆号擷取最新Linux技術分享。