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

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

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

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

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

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

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

模型架構

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

圖：模型架構和增強模塊可視化

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

ML 管線(xiàn)

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

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

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

結果呈現

點(diǎn)擊圖片，查看弱光增強前后對比

<< 滑動(dòng)解鎖更多示例 >>

性能優(yōu)化

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

量化

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

模型推理時(shí)間

圖：模型推理時(shí)間 Pixel 7

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

安卓性能提示 API

想要重復對演示應用進(jìn)行基準測試，就必須要用 ADB 命令來(lái)開(kāi)啟安卓固定性能模式。這是因為每次測試時(shí)，CPU 的頻率可能會(huì )變來(lái)變去，ADB 命令可以固定 CPU 頻率。我們發(fā)現，在使用這種固定性能模式后，幀時(shí)間減少了。但這也讓?xiě)瞄_(kāi)發(fā)者左右為難，因為他們沒(méi)法控制 CPU 頻率，但又不能指望終端用戶(hù)會(huì )使用 ADB。不過(guò)，安卓性能提示 API 可以解決這個(gè)問(wèn)題。該 API 主要用于游戲，其工作原理是設定目標幀時(shí)間并將該指標報告回安卓系統，由安卓系統調整時(shí)鐘頻率，嘗試達到該目標。這使得幀時(shí)間得到了與固定性能模式相當的良好改善。

幀率

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

進(jìn)一步的模型訓練

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

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

結論

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