-01-
AI編解碼的意義與挑戰
以上是高通總結的關于AI編解碼優勢的一張示意圖。其中相當多的優勢來自于端到端優化這一特性,它的壓縮率比較好,可以對任意分布的數據做專門優化,可以針對任意更符合主觀質量的、更符合下游任務的損失函數進行優化。
另一方面,由于它是用神經網絡實現的,那便可以在人工智能的大潮中進行復用,例如復用各種推理硬件和算法。還有一個好處是權重比較容易更新,傳統編解碼算法在做成硬件后很難進行修改,但對于基于神經網絡的AI編解碼算法,它的權重是可以進行修改的,這是一個很大的優勢。
接下來介紹個人認為對AI編解碼器較為重要的六個評價維度,第一是率失真性能和主觀質量,它主要和壓縮率有關;第二是復雜度,它與延時、計算量和顯存的要求還有功耗、吞吐率等因素有關;第三是跨平臺解碼,在手機、CPU、GPU上互相編解碼不應該出錯;第四是對下游AI任務訓練和推理的影響,對測試或推理的影響類似于現在比較熱門的面向機器視覺的編解碼,在訓練方面大家會比較關注用AI壓縮的數據是否會對神經網絡的訓練效果產生影響;第五是泛化與特化能力,有時我們希望它可以泛化,使用同一個模型可以壓縮不同數據,有時我們希望它可以特異化,例如在壓縮遙感或者醫學類數據時,可以構造專門的模型使對應數據的壓縮率更高;最后則是轉碼穩定性,這也是一個有意思的問題,包括傳統算法和AI算法間的互編互解,例如將JPEG二次壓縮再解碼,過程中是否有性能損失。
-02-
提升RD性能與解碼速度
接下來介紹如何提升RD性能和解碼速度,圖中的白色文字為前人所做的一些早期經典研究成果。第一篇為紐約大學的論文,模型成果第一次在PSNR上超過JPEG 2000。第三篇論文的模型成果首次在PSNR上超過BPG,BPG對應H.265的幀內壓縮技術,但解碼速度下降了約60倍。
我們的工作自此開始展開,具體表述為下面的黃字部分。首先我們于2021年構造了棋盤格上下文模型,消除了60倍復雜度,又于2022年對模型進行了一些新的改進。
以上是紐約大學在ICLR上發表的論文,它是一個簡單的變分自編碼器(Variational Auto-Encoders)。它首先將圖像變到一個特征域進行概率估計,然后用熵編碼對特征進行壓縮。該模型是2017年的研究成果,當時已經超過了JPEG 2000。
論文作者在加入谷歌后構造了一個新的尺度超先驗模型(Scale Hyperprior)。它的思想是:在壓縮神經網絡特征時我們實際使用的是獨立假設,由于它是一個很高維的向量,所以不能使用聯合概率密度分布。該模型通過引入先驗Z使得要壓縮的特征Y變為條件獨立,因而可以使用概率密度估計和熵編碼算法將特征Y壓縮得更小。該模型中特征Y和先驗Z的碼率之和比2017年論文中Y的碼率更小,這是通過引入更精準的數學建模帶來的提升。
谷歌接著又對模型做了進一步改進,即前面提到的空間自回歸模型(autoregressive),它的思想是:在壓縮碼字Y時我們并不能做到獨立壓縮,如在壓縮左側方格圖中的紅色像素時,需一并考慮周邊的黃色像素結果,這與傳統算法中的幀內預測相似,被稱為自回歸。模型效果如右圖所示,其成果第一次超過了BPG,非常有代表性。
接下來介紹我們的算法,剛才提到的自回歸模型要逐個像素串行解碼,速度相當慢。我們在此基礎上提出右圖所示的棋盤格模型。例如在第一次解碼時,我們使用GPU來并行解碼右側方格圖中所有藍色和黃色的點,第二次解碼紅色和白色點,由此可通過兩次解碼解出整張圖。按照谷歌的原方法,串行復雜度為N^2,使用棋盤格模型,串行復雜度僅為2。
以上為測速結果,在典型模型上棋盤格模型能加速40到50倍,在變換網絡較大的模型上能加速18到20倍。該研究成果發表于CVPR 2021。
在2022年CVPR中,我們對原論文進行了改進,將谷歌的自回歸模型和棋盤格模型進行了結合,成為了一種既有通道自回歸、也有空間上下文的并行上下文模型,該模型取得了非常好的效果。
以上為對模型所做的另一項改進,即將通道劃分改為非均勻劃分。谷歌原模型每個通道分組都為32,需要的分組較多。我們發現實際上其中信息的分布是不對稱的,只有少數信息分組比較重要。為了提高速度,我們將非重要分組進行了合并,在提高速度的同時甚至提高了壓縮率。
以上是我們新模型的結果,顯示為紅色,最終要優于VVC,并且與之前的一些網絡速度相近。圖表中橫軸為解碼速度,縱軸為壓縮率。
上表可以更加直觀體現論文研究成果的貢獻,這實際是端到端圖像壓縮領域的一個里程碑式的進展??梢钥吹剿性贐D-Rate上顯示為負數的優于VVC的方法,其解碼速度都要大于1000,解碼一張圖大約要大于1秒。早期論文的解碼速度都在20到50毫秒。我們的論文首次實現了BD-Rate在約-7%時,解碼速度仍只有50毫秒,這是相當快的速度。
此后我們對VAE框架進行了進一步研究,VAE框架中存在均攤變分推理現象,利用這個現象引入半均攤變分推理,可以在編碼時對碼字進行一定更新并實現很多靈活的控制。如可以實現連續變碼率,使用同一個解碼器可以解不同碼率的圖像。可以實現任意ROI編碼,還可以去權衡不同的質量評價指標,比如PSNR和LPIPS,大家知道這些指標和主觀質量的相關性各有不同,權衡不同指標等于在權衡解碼圖像的不同特點。我們將相關方法也擴展到了端到端視頻編解碼上,相關成果發表在ICML 2023。
這是我們發在NeurIPS 2022上的一篇論文,端到端圖像壓縮領域除了調整網絡結構外,其實會有比較深入的理論背景,我們對理論進行了更深入的探索,將單樣本采樣改為多樣本采樣。造成的結果是訓練速度會變慢一些,但壓縮和解壓縮速度完全不受影響。利用這個技術我們取得了一個較明顯的壓縮率的提升。
-03-
提升主觀質量
下面介紹一下如何提升主觀質量。一般編解碼器在進行設計時,優化指標為PSNR、 SSIM或VMAF。但它們與主觀質量的差距都較大,實際上也不存在與主觀質量絕對一致的數學指標,這給我們造成了很大困難,一般我們會選擇同時優化多個數學指標來使最終的主觀質量變好。
在端到端圖像壓縮領域一般會有兩個較常見的數學指標:PSNR和MS-SSIM,我們希望設計一種模型,在其訓練后使兩種指標都較高,最后的主觀質量更好。
通過實驗發現,我們能夠做到在PSNR不下降的情況下盡可能提升SSIM指標,使主觀質量得到提升。
以上為展示效果圖,左上角為原圖,最右側為VVC結果。中間以SSIM標記的是使用MS-SSIM做損失函數訓練的模型所優化的結果,最終會出現色偏。用MSE訓練的結果在碼率較低的位置(如圖中水的位置)會損失紋理。使用我們的方法訓練的兩個模型可以較好的平衡PSNR和SSIM,不發生色偏和紋理損失。
在上圖中情況是相同的,尤其在草地部分,MSE結果草地會較糊,SSIM結果草地顏色會出現偏差,而經過改進后模型的效果是較好的。
谷歌也相當重視提升主觀質量這方面工作,在每年的CVPR上都會組織圖像壓縮競賽,該賽事擁有比較完善的主觀質量評估流程,通過眾包的方式請很多人來看圖,會規定圖片的分辨率和與屏幕的不同距離來評估解壓圖像的主觀質量,競賽的組織者一般是國際上比較知名的一些廠商。
我們也參加了該賽事,使用的方法就是同時優化多個數學指標,第一個是感知損失,第二是重建損失(例如PSNR或SSIM),第三是對抗損失,第四是風格損失,其中三個損失函數都和深度學習有關。我們將四個損失函數以一種特定系數去進行組合,最后優化出來的模型在各種評價指標上都是最佳的。
我們以該模型解壓的圖像參賽,最終獲得了所有碼點的第一名。
這是谷歌官網對所有方法的評測結果,首先我們關注一下主觀質量和數學指標間的關系,例如PSNR,我們可以看到像avif、VVC和HEVC方法的PSNR都較高,但對應的ELO列(人眼看圖的主觀質量)都不太好,印證了PSNR和SSIM不代表主觀質量這一結論。
另一方面,可以看到我們方法的解碼速度,它是用PyTorch編寫的,解碼速度達到了460,VVC在谷歌測試中是493。我們的模型主觀質量比VVC更好,解碼速度還要更快,但比avif在解碼速度上要慢一些。
不過需要注意的是,它只是參考軟件層面的對比,因為這里的VVC是使用VTM進行測試,AI方法一般是使用PyTorch進行測試,所有方法都沒有進行工業級的性能優化。
-04-
跨平臺編碼
下面講一個看起來比較偏,但實際對編解碼比較重要的問題,即跨平臺解碼。
這個問題主要來自于熵編碼,我們進行熵編碼或算數編碼時每個碼字的概率不能出錯,如果有任何微小的錯誤,后面的碼點便解不出來。如上圖所示,前面一直在正常解碼,一旦算到某處一個像素不對,那么后面所有的解碼都會出錯,概率錯了所有碼點都解不出來。假如在壓縮圖片時原先使用的硬件丟失,那么圖片數據也會永久消失。
這個問題的解決辦法是使圖片的編解碼過程,尤其是熵編碼的概率計算過程不管在何種硬件上(如CPU、GPU還是不同型號的NPU或DSP)計算結果都完全一致。據我所知,唯一的方法便是使用完全整數計算來實現。
首先針對概率計算有關的所有神經網絡,我們均使用全整數推理,同時以使用查找表而非直接計算的方式來進行概率計算。這樣所有過程都是用整數來實現的,可以確保編解碼具備跨平臺條件。以上工作比較細節,所以我們公開了一個比較詳細的技術報告,有興趣實現技術落地的可以參考。
上圖為性能對比,我們將新方法和之前類似的方法進行了對比,我們是首先在有上下文的圖像壓縮模型上實現不掉點的整數推理,而且我們的量化方案比較標準,沒用對激活值進行逐個通道的分組量化,適用于常見的GPU和NPU。
-05-
優化延時與吞吐
最后關于有損壓縮方面介紹一下優化延時與吞吐。優化一個AI編解碼器的速度主要包括兩部分工作:一是優化神經網絡推理的延時,另一個是優化熵編碼的延時。最后整個系統要做一個代碼層面或者軟件工程層面的優化。
首先看一下神經網絡的優化,它的方法都是比較標準的。其中一個是神經網絡架構搜索(NAS),一個是模型量化。NAS提升顯得較小是因為這些神經網絡已經經過了人工優化。左下角是它的CPU和GPU顯存占用情況。由于針對此模型我們主要調整它的主觀質量,所以它的PSNR會弱一些??梢钥吹皆撃P偷腜SNR比avif低,但SSIM值很高。
在本頁我們想回答一個問題,神經網絡圖像壓縮和傳統方法相比性能有何差距?我們希望進行相對公平的對比。首先要使用相同的硬件,例如都使用CPU,以編解碼延時來衡量,神經網絡圖像壓縮和JPEG XL和avif在性能上是接近的,如果我們將avif和JPEG XL作為工業應用的發展方向,顯然神經網絡圖像壓縮也可以作為一個發展方向。
如果我們看GPU,可以發現在相同的輸入條件下,神經網絡方法和JPEG-turbo相比,它在編解碼上并沒有很大的劣勢,解碼還要快一些。
但這里有一個關于線程的小問題,從JPEG官網來看,它的測試也會有一些問題,就是到底采用什么樣的線程數來測試這些模型,使用不同的線程數測試結果也不同。我們采用了單線程進行測試,對于傳統方法采用的是默認配置,使用多線程在測吞吐時會有更大優勢,具體可以參考我們的技術報告。這個結果可以作為一個參考,實際使用中應該結合運行環境進一步對線程資源進行適配。
-06-
應用拓展:JPEG無損壓縮
剛才介紹了有損壓縮,接下來分享一個比較有意思的應用,即無損壓縮。假如端到端有損壓縮方法實際落地要等很長時間,那么我們現在已經有了非常多的JPEG圖像,能不能考慮使用AI方法對這些圖像進行無損壓縮。
其實已經有人注意到這個問題,但使用的不是神經網絡,例如JPEG XL或者Dropbox的Lepton,Lepton目前已經被Dropbox使用了很久,它的思想是在進行云存儲時對JPEG圖像進行無損二次壓縮,在不改變用戶數據的同時極大壓縮存儲量。所以它整個生態不會受到任何影響,在網絡上還是以JPEG格式進行傳遞,但是云存儲使用的是二次壓縮后的碼流。
就此我們在CVPR上也發表了一篇文章,專門介紹如何使用神經網絡來無損壓縮JPEG圖像。在這里對研究內容做一個簡單介紹,首先我們將JPEG圖像的三個分量(YUV或YCbCr)整合成一個引變量Z,整合后的Z涵蓋了YUV之間的關聯,在壓縮YUV三個分量前使用Z來作為先驗。Y叫作亮度分量,由于它包含的信息最多,所以我們為這個分量專門設計一個模型叫作MLCC,它的結構如下。
MLCC是一種比較復雜的并行自回歸模型。左圖中為JPEG的dct系數,我們對它進行重新排列,將相同頻率排到神經網絡的相同通道上得到右圖,并將右圖中三維長方體的行列進行展開,按照類似自回歸的方式進行條件建模。它看起來很復雜,但實際執行的速度還是很快的。
以上為測速結果,Dropbox在大規模使用的Lepton,其無損壓縮率是20%左右,編碼解碼時間分別是0.239和0127。JPEG XL稍差,但編碼稍微快一些。作為傳統算法中號稱最強的無損壓縮器CMIX,它的壓縮率能達到23%或27%,但編碼和解碼時間都很長,要150秒。該軟件是壓縮從業者為了探索無損壓縮極限構造出的一套十分復雜的模型,它的算力消耗和內存消耗都非常非常大。
我們的神經網絡方法壓縮率可以達到29%或30%以上,優于CMIX,使用PyTorch在GPU上的編解碼時間大約為1秒,比Lepton要慢約十倍。
不過我們要知道Lepton不是一個參考軟件,而是一個工業軟件,我們的算法作為學術上的參考軟件,兩者gap并不大。
我們也將神經網絡方法和之前做無損壓縮的一些網絡結構進行了對比,可以看到本方法的編解碼時間實際是比較短的。
小結
下面進行一個小結。經過剛才提到的一些優化,AI編解碼方法整體上在一些場景下是可以使用的,但在我剛才提到的六個維度上還有一些持續的挑戰。
本次分享的內容來自這些已經發表或公開的論文和技術報告,感謝商湯科技和清華大學的合作者。
審核編輯:劉清
評論
查看更多