使用移動端神經(jīng)網(wǎng)絡(luò)實現(xiàn)實時弱光視頻增強

作者：Arm 戰(zhàn)略與生態(tài)部工程師 Ayaan Masood

作為通訊工具，視頻會議幾乎隨處可見，尤其適用于遠程辦公和社交互動。但其使用體驗并非總是簡單直接、即開即用，可能需要進行調(diào)整，確保音頻視頻設(shè)置良好。其中，照明便是一個難以把握的因素。在會議中，光線充足的視頻畫面會顯得得體大方，而糟糕的照明條件則會顯得不夠?qū)I(yè)，還會分散其他與會者的注意力。有時，改變光照情況并不可行，特別是在光線昏暗的冬季或照明不足的地點。

在本文中，我們將探討如何構(gòu)建一個演示移動端應(yīng)用，以解決弱光條件下的視頻照明問題。我們將介紹支持該應(yīng)用的神經(jīng)網(wǎng)絡(luò)模型及其機器學(xué)習(xí) (ML) 管線、性能優(yōu)化等。

找到合適的神經(jīng)網(wǎng)絡(luò)

我們采用基于神經(jīng)網(wǎng)絡(luò)的解決方案改善視頻照明。因此，這項工作的核心在于找到適合當(dāng)前任務(wù)的神經(jīng)網(wǎng)絡(luò)。目前，市面上有很多出色的開源模型可供使用，而為本項目找到合適的候選模型至關(guān)重要。我們在評估模型時主要關(guān)注以下三項要求：性能良好、照明質(zhì)量出色、視頻處理表現(xiàn)優(yōu)異。

我們的目標是實現(xiàn)移動端實時推理，這意味著，每幀處理時間只有嚴格控制在 33 毫秒內(nèi)，才能實現(xiàn)每秒 30 幀的流暢播放。其中包括預(yù)處理/后處理步驟以及神經(jīng)網(wǎng)絡(luò)運行所需的時間。視頻畫質(zhì)增強是另一項重要標準。該模型可智能化增強暗像、還原細節(jié)，確保視頻幀之間暫時保持連貫，避免畫面閃爍。

模型架構(gòu)

所選模型來自 2021 年研究論文《用于弱光圖像/視頻增強的語義引導(dǎo)零樣本學(xué)習(xí)》 ^[1] 。在包含混合曝光和照明條件的復(fù)雜數(shù)據(jù)集上進行測試時，該模型的弱光增強質(zhì)量非常出色。暗像中不清晰的細節(jié)和結(jié)構(gòu)突然變得清晰起來。該模型的另一個優(yōu)點是尺寸極小，僅有一萬個網(wǎng)絡(luò)參數(shù)，因此推理速度很快。在模型架構(gòu)方面，輸入的圖像張量被縮小并傳遞到卷積層棧。這些層會預(yù)測逐像素增強因子，然后在模型后處理模塊中，將增強因子以乘法方式應(yīng)用到原始圖像像素上，從而生成增強的圖像。

圖：模型架構(gòu)和增強模塊可視化

值得一提的是，訓(xùn)練采用包含兩千張合成圖像的小型數(shù)據(jù)集，這就表明，如果數(shù)據(jù)集規(guī)模更大，模型就還有更大提升空間。我們增強了真實參照圖像，生成一系列統(tǒng)一曝光值，從而調(diào)整圖像明暗。該模型的訓(xùn)練不受監(jiān)督，無需標簽，而是通過一個指導(dǎo)訓(xùn)練的損失函數(shù)來學(xué)習(xí)如何增強弱光圖像。這個損失函數(shù)是多個獨立損失的組合，它們分別負責(zé)圖像的各個方面，如顏色、亮度和語義信息。

ML 管線

該應(yīng)用的 ML 管線從輸入幀開始，根據(jù)模型輸入張量要求加以處理，然后推理并向用戶顯示輸出。攔截相機幀進行推理是 CameraX 庫的內(nèi)置功能，通過安卓的“圖像分析”API 實現(xiàn)。

我們采用了兩種不同的 ML 推理引擎：ONNX  runtime和 TensorFlow Lite。將 Pytorch 模型導(dǎo)出為 ONNX 模型是 Pytorch 庫自帶的功能，而導(dǎo)出到 TensorFlow Lite 則困難得多。該模型的有效導(dǎo)出器為 Nobuco，其工作機制是先創(chuàng)建可轉(zhuǎn)換為 TFLite 的 Keras 模型。

模型推理產(chǎn)生的輸出格式取決于 ML 運行時。如果采用 ONNX，則為 NCHW（即數(shù)量、通道、高度、寬度）；如果采用 TFLite，則為 NHWC，其中通道排在末尾。這影響了后處理步驟的完成方式，即將整數(shù) RGB 值的輸出緩沖區(qū)解包，以創(chuàng)建終末位圖，并顯示在屏幕上。

結(jié)果呈現(xiàn)

點擊圖片，查看弱光增強前后對比

<< 滑動解鎖更多示例 >>

性能優(yōu)化

在 Kotlin 中將 RGBA 位圖轉(zhuǎn)換為 RGB 的計算成本很高。當(dāng)性能預(yù)算限制在 33 毫秒以內(nèi)時，僅轉(zhuǎn)換過程就需要花費幾十毫秒。為了加快速度，我們使用了 C++ 語言，并對編譯器進行了全面優(yōu)化。但是，要從 Kotlin 代碼調(diào)用 C++ 代碼，就得用到 Java 原生接口 (JNI)。通過 JNI 傳遞一個 3x512x512 大小的浮點數(shù)緩沖區(qū)成本很高，因為必須復(fù)制兩次，先從 Kotlin 復(fù)制到 C++，處理完后再復(fù)制回 Kotlin。為了解決這個問題，我們使用了 Java 直接緩沖區(qū)。傳統(tǒng)緩沖區(qū)的內(nèi)存由安卓運行時在堆上分配，C++ 不容易訪問。而直接緩沖區(qū)必須按照系統(tǒng)的正確字節(jié)順序分配，而一旦分配好了，內(nèi)存就能以操作系統(tǒng)和 C++ 易于訪問的方式分配。這樣，我們就省去了復(fù)制到 JNI 的時間，并能夠利用高度優(yōu)化的 C++ 代碼。

量化

該模型使用量化技術(shù)進行優(yōu)化。量化就是使用較低精度表示神經(jīng)網(wǎng)絡(luò)的權(quán)重和激活值，從而在略微犧牲模型質(zhì)量的情況下提高推理速度。量化用的數(shù)據(jù)類型通常比較小，例如 INT8，它占用的空間只有原來 32 位浮點數(shù) (FP32) 的四分之一。模型量化有兩種：動態(tài)量化和靜態(tài)量化。動態(tài)量化僅量化模型權(quán)重，并在運行時，針對激活值確定量化參數(shù)。靜態(tài)量化則事先使用代表性數(shù)據(jù)集量化權(quán)重和激活值，所以在推理方面更快。對于這個模型，靜態(tài)量化提高了推理速度，而輸出照明增強效果則略暗一些，這樣的取舍是值得的。

模型推理時間

圖：模型推理時間 Pixel 7

上圖比較了使用 ONNX  runtime 和 TensorFlow Lite 對弱光增強模型進行模型推理的時間。在 Pixel 7 上的分辨率為 512x512。 我們先將 ONNX  runtime 作為推理引擎。在 CPU 上運行時，F(xiàn)P32 的模型推理時間為 40 毫秒。當(dāng)量化到 INT8 時，這一時間縮短到 32 毫秒。性能原本有望獲得大幅提升，但使用可視化工具 Netron 分析模型文件后發(fā)現(xiàn)，模型圖中增加了額外的量化/反量化算子，從而增加了計算開銷。使用 XNNPack 和 INT8 模型的 TensorFlow Lite 在 CPU 上的表現(xiàn)稍慢于 ONNX runtime，推理時間接近 70 毫秒。不過，TensorFlow Lite 在使用 GPU 處理時，超越了之前所有推理引擎和模型類型的組合。對于 512x512 輸入圖像的推理，TensorFlow Lite 僅需 11 毫秒，因此我們選擇它作為運行模型的后端，以實現(xiàn)實時照明增強。

安卓性能提示 API

想要重復(fù)對演示應(yīng)用進行基準測試，就必須要用 ADB 命令來開啟安卓固定性能模式。這是因為每次測試時，CPU 的頻率可能會變來變?nèi)?，ADB 命令可以固定 CPU 頻率。我們發(fā)現(xiàn)，在使用這種固定性能模式后，幀時間減少了。但這也讓應(yīng)用開發(fā)者左右為難，因為他們沒法控制 CPU 頻率，但又不能指望終端用戶會使用 ADB。不過，安卓性能提示 API 可以解決這個問題。該 API 主要用于游戲，其工作原理是設(shè)定目標幀時間并將該指標報告回安卓系統(tǒng)，由安卓系統(tǒng)調(diào)整時鐘頻率，嘗試達到該目標。這使得幀時間得到了與固定性能模式相當(dāng)?shù)牧己酶纳啤?/p>

幀率

當(dāng)弱光增強功能打開時，應(yīng)用會顯示幀時序。盡管幀率約為 37 FPS，但攝像頭幀速率會根據(jù)硬件和照明情況受到限制（在極弱光條件下，安卓系統(tǒng)會降低相機 FPS 以提高亮度）。在 Pixel 7 上，向用戶顯示的幀速率（默認的安卓攝像頭 API）上限為 30 FPS。更快的推理并不會帶來更好的用戶體驗，因此保留了一個 7 FPS 的性能緩沖。

進一步的模型訓(xùn)練

盡管原始模型在暗場景下的照明增強效果相當(dāng)不錯，但圖像有時會出現(xiàn)過度曝光的情況。為了解決這個問題，我們采用了一個包含兩萬張合成圖像的大型數(shù)據(jù)集進行訓(xùn)練，此前的研究論文僅使用了兩千張圖像。

不過，基于大型數(shù)據(jù)集的訓(xùn)練時間會變長。經(jīng)過調(diào)查，性能下降的原因是每批量八張圖像超出了 GPU VRAM 容量，并溢出到系統(tǒng)內(nèi)存中。為了在不增加 VRAM 使用量的情況下提高有效批次大小，我們采用了梯度累積技術(shù)，無需針對每個批量計算梯度，而是累積多個批量后再計算梯度。在我們的案例中，我們可以使用的批量上限是六張圖像，而采用梯度累積技術(shù)后，我們能夠使用的批量是 60 張圖像。

結(jié)論

我們在本文中展示了一個可運行的演示移動端應(yīng)用，用于實時改善移動端視頻的照明效果。基于 Arm 平臺優(yōu)化并運行 ML 模型的過程非常順暢，得益于量化和各種推理引擎等技術(shù)的運用，該模型能夠在 33 毫秒的嚴苛性能限制下順利運行。