概述
华为诺亚方舟实验室AI基础理论团队近期提出了一种基于深度生成模型的实用图像无损压缩框架PILC。该框架在模型实现上结合了自回归模型以及自编码模型,将图片建模转化为残差建模,同时针对编码器进行了优化及GPU适配,使得整个框架均运行在GPU上,大大提升压缩解压吞吐率。图1显示了PILC与现有方法在压缩率和压缩吞吐率上的比较。
图1 PILC与现有无损压缩算法比较
可以看出,PILC在压缩率上超过传统方法及深度学习方法L3C,同时在压缩吞吐率上大幅超过现有的基于深度生成模型的方法,将压缩解压吞吐率提高到了约200MB/s,为AI无损压缩的商业化提供了新的思路。相关工作已被CVPR 2022接收,论文标题是: PILC: Practical Image Lossless Compression with an End-to-end GPU Oriented Neural Framework。
文章链接:
研究背景
无损压缩的基本出发点是用更短的码编码出现频率更高的符号,香农信源编码定理给出了数据无损压缩的下限就是数据的信息熵,因此只要我们对数据的概率分布预估的越准,就可以逼近数据压缩的下限,达到更高的压缩率。近年来,基于深度生成模型的图像无损压缩技术快速发展,正是利用了深度生成模型在概率估计上取得的成功,自回归模型、自编码模型、流模型等生成模型均被应用于图像无损压缩,但是目前的工作在网络推理速度、编码器、数据传输上存在瓶颈,导致压缩解压吞吐率很低,难以实用。
方案描述
1、 组合模型进行残差建模
在自回归模型、自编码模型、流模型三类生成模型中,自编码模型拥有最快的推理速度。但是之前的工作表明,自编码器对局部特征并不能很好的进行捕获,因此概率建模准确性并不是最高的。因此我们考虑利用一个局部的自回归模型建模图像局部特征,然后用自编码器建模全局特征。
图2 总体框架
总体框架如图2所示。压缩时,自回归模型首先对输入图片进行重建,得到图片残差,VQ-VAE模型建模残差的概率分布,之后编码器会对残差进行压缩,同时我们还需要压缩VQ-VAE中的向量序号,以便再解压时得到残差的概率分布。解压时,首先解压向量序号,接着得到残差的概率分布,之后根据概率分布解压出残差,自回归模型会将残差恢复到原图。得益于局部自回归的设计和轻量化的VQ-VAE,组合模型的整体参数很少,模型推理吞吐率能够达到276MB/s。
2、 自回归模型并行解压策略
图3 自回归策略及并行设计
为了达到较高的吞吐率,自回归模型的设计应该尽可能小,图3左侧图片展示了我们三路自回归模型的设计。对于R通道的每一个点,它的预测仅依赖于左侧、上侧、左上的三个点,对于G、B通道的每一个点,我们还格外考虑了跨通道的信息。
传统自回归模型在解压时每一个点必须依赖之前的所有点进行解码。而我们的设计策略使得解压能够进行并行,具体来说对于R通道的点,根据预测规则,一条反对角线上的所有点可以同时完成预测,对于G、B通道的点,每一列直线上的所有点可以同时完成预测。实验验证,模型解码时有并行策略和无并行策略的吞吐率分别是382.5MB/s和48.5MB/s。
3、 编码器设计
编码器也是AI压缩吞吐率的一大瓶颈。传统非AI算法多采用静态概率建模,可用静态编码器;而AI算法由于采用多种条件概率,只能采用动态编码器。相对于静态,动态编码器更为灵活、通用,但吞吐率通常有数十、数百倍的差距。在PILC工作中,我们设计了一种高效动态编码器ANS-AI。其借鉴了静态熵编码中利用建表实现快速存取计算结果的思路,并设计了一系列的近似分布运算,从而避免了表格内存占用过大的问题。实验结果表明,ANS-AI可在压缩率损失4%以内的前提下,吞吐率与静态编码器相当。
实验结果
1、 跟现有无损压缩算法的对比
表1 压缩性能对比
表1我们比较了PILC与传统方法及基于自编码模型的L3C在压缩效果与压缩解压速度的差异,我们的框架取得了最快的压缩解压吞吐率,达到约200MB/s。在低分辨率图像上的压缩率很有竞争力,在高分辨率图像上的压缩率也依然领先传统方法。
2、 每阶段压缩解压吞吐率
表2 每阶段压缩解压吞吐率及单张图处理时间
表2展示了PILC压缩解压图片全流程每个阶段的吞吐率及单张图片的处理时间,同时我们也给出了编码器在CPU和GPU上运行的对比。在压缩阶段,编码器在CPU和GPU上的运行区别不大,因为主要流程是一样的。而在解压时,编码器在CPU上运行会增加两次额外的GPU-RAM之间的传输,这主要是因为VQ-VAE模型在GPU上推理出残差概率分布后需要传到CPU的编码器,编码器根据概率分布解码出残差,之后残差需要再传送到GPU上进行原图的重建。可以看出,在GPU上运行Coder会减少一定的数据传输次数,对吞吐率的提升有帮助。
3、 消融实验
表3展示了不同模型结构在压缩率上的效果。可以看出VQ-VAE的效果要好于AE,表明向量索引结合码本可以包含具有更多多样性特征,有助于概率预测的准确性。PILC中的AR+VQ-VAE模型在压缩率上取得了最好的效果,这表明了预测残差设计的有效性。
表4展示了自回归模型使用不同的感受野的对比,这里感受野指的是当前点的预测需要依赖之前多少个点。可以看出三路自回归加并行设计能够取得最高的吞吐率,进一步增大感受野带来的提升是有限的。
总结
本文提出了一种基于深度生成模型的实用图像无损压缩框架PILC,它能够在压缩率超过传统方法及一些基于生成模型的方法的同时,取得较高的压缩解压吞吐率,对基于深度生成模型的无损压缩技术进一步商业化提供了新的思路。
作者:诺亚方舟实验室
文章来源:诺亚实验室
推荐阅读
更多嵌入式AI相关技术干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。