一、文章摘要
图像重缩放是一种常用的双向操作,它首先将高分辨率图像缩小以适应各种显示器或存储和带宽友好,然后将相应的低分辨率图像放大以恢复原始分辨率或放大图像中的细节。然而,非单射下采样映射丢弃了高频内容,导致逆恢复任务存在不适定问题。这可以抽象为具有信息损失的一般图像退化-恢复问题。在这项工作中,我们提出了一个新的可逆框架来处理这一普遍问题,该框架从一个新的角度来模拟双向退化和恢复,即可逆双射变换。该框架的可逆性使其能够以分布的形式对退化前的信息损失进行建模,从而缓解恢复后的不适定问题。具体来说,我们开发了可逆模型来生成有效的退化图像,同时在正向退化过程中将丢失内容的分布转化为潜变量的固定分布。然后,通过对生成的退化图像与随机绘制的潜在变量应用逆变换,使恢复变得易于处理。我们从图像的重缩放开始,实例化模型为可逆的重缩放网络,可以很容易地扩展到类似的脱色-着色任务。我们进一步建议将可逆框架与现有的退化方法(如图像压缩)结合起来,以获得更广泛的应用。实验结果表明,与现有方法相比,我们的模型在从缩小和脱色图像进行放大和彩色重建的定量和定性评估以及图像压缩的速率失真方面都有显著改进。
二、文章提出的方法
2.1 图像重缩放模型
对于具体的模型,我们在这一节从图像重缩放开始。我们开发了可逆重缩放网络(IRN)作为我们图像重缩放可逆建模框架的实例化模型,并描述了IRN的具体可逆结构和训练方法。我们还在算法1,2中给出了IRN模型中的下采样和升尺度算法作为示例,以更好地演示我们的可逆框架的输入、输出和过程。注意,在实践中,HR图像x和LR图像y将被量化为8位表示,如3.3.1节所述。我们在算法描述中省略了这个细节,并将定义域视为R。
2.1.1 可逆结构
图2说明了我们提出的IRN的架构,它基本上由一个Haar变换和几个InvBlock组成的堆叠降尺度模块组成。每个降尺度模块将空间分辨率降低2倍。如果每个组件都是可逆的,那么整个体系结构就是可逆的。
图2 将我们的可逆重缩放网络(IRN)作为可逆建模框架的实例化模型的说明。可逆体系结构由降尺度模块组成,其中InvBlock在Haar变换后堆叠。每个降尺度模块将空间分辨率降低2×。ρ的exp(·)被省略
Haar变换 在每个降尺度模块中,首先应用Haar变换,使模型具有一定的感应偏差,用于分割低频和高频内容,在图像下采样过程中,低频和高频内容分别被近似保留和丢失。Haar变换是一种可逆小波变换,它将输入分解为低通表示和三个方向的高频系数。具体来说,给定输入高H、宽W、通道C的原始图像或特征图,生成一个形状张量(1/2 H,1/2 W,4C),其中第一个C片是等价于双线性插值下采样的低通表示,其余三组C片分别对应于垂直、水平和对角线方向的高频残差。通过Haar变换,模型能有效分离低频和高频信息,有利于y的后续生成和xH到z的变换,Haar变换后的空间分辨率降低了2倍。
InvBlock InvBlock是目标可逆转换的主要组件。考虑到输入已经被Haar变换分成低频和高频成分,我们引入了基于方程(1,2,3)中描述的耦合层架构的InvBlocks,其两个分支(即式(1)中hl1和hl2的分割)分别对应这两个分量。该转换将进一步优化输入表示,以生成合适的LR图像,以及针对丢失信息的独立且适当分布的潜在表示。对于详细的计算,考虑到图像缩放任务中快捷连接的重要性,我们对低频部分hl1采用加性变换(Eq.1),对高频部分hl2采用增强仿射变换(Eq.3)来增强模型容量。这也使模型具有一定的感应偏置,用于生成低频部分直接通过的LR图像,可以稳定IRN的训练。InvBlock架构的细节如图2所示,除了这里省略了函数ρ后的exp(·)操作。
我们使用一个密集连接的卷积块来参数化变换函数φ(·),η(·),ρ(·),该卷积块已在中证明了其对图像缩放任务的有效性。为了避免由于exp(·)函数引起的数值爆炸,我们在函数ρ(·)之后使用了居中的S形函数和标度项。
量化 我们的模型的输出是浮点值,而常见的图像格式(如RGB)则被量化为8位表示。为了实现存储兼容性,我们在生成的LR图像上采用舍入操作作为量化模块。量化后的LR图像保存为PNG格式,用于缩放。然而,量化的不可微特性给反向传播训练带来了挑战。为了克服障碍,我们应用直通估计器来计算量化模块的梯度。为了简单起见,下面省略了量化的符号。
2.1.2 规模灵活高效的实施
还可以进一步改进架构,使IRN适应更大的规模或更高的计算效率。具体来说,我们将引入可学习的下采样模块和计算效率的改进,以实现规模灵活和高效的实现。
可学习的下采样 虽然Haar小波变换能够很好地进行高频和低频内容的下采样和分割,但叠加多次变换只能以2的幂来重缩放图像。这在很大程度上限制了我们模型的缩放范围。为了实现更多的尺度,例如3×,我们建议利用一个可学习的下采样层来取代架构中的Haar变换。它由一个压缩运算和一个1 × 1可逆卷积组成。
如图3所示,压缩操作通过将空间元素压缩到通道中,降低了特定尺度N的空间分辨率。然后,利用1 × 1可逆卷积对压缩后的N × N个元素进行InvBlocks前的变换。1 × 1可逆卷积首次在GLOW (Kingma & Dhariwal,2018)中提出用于信道排列。不同于它们的目的,我们希望它能学会在任意尺度下分割低频和高频内容,并更好地适应以下InvBlocks。Haar变换可以被视为2倍尺度下这个降尺度模块的一个特例,因为它提供了一个固定的而不是可学习的先验。对于该模块,我们通过设置1 × 1可逆卷积的参数,使变换后的第一个通道为N × N个元素的平均值,而其他通道为恒等变换,使其具有可逆性,从而提供初始化时提取低频的先验。
图3 可学习的下采样模块的说明(2×示例)。它包括将空间分辨率降尺度N倍的压缩操作和对压缩后的N × N个元素进行1× 1可逆卷积变换
我们将具有可学习下采样的IRN模型表示为IRNLD。
分数缩放因子 在实际应用中,会有分数缩放因子。我们可以将IRN与传统插值方法结合起来进行处理。具体而言,对于缩放因子s1,我们选择缩放因子s2 = [s1]的IRN,并在将HR图像传递给IRN之前和之后分别通过缩放s2/s1和s1/s2进行插值(例如双三次)重缩放HR图像。最近的工作也证明了这一点(Xing等,2022)。
提高计算效率 我们注意到,堆叠多个降尺度模块(包含一个下采样模块和多个InvBlocks)的架构在计算过程中遭受了大量增加的FLOP。这是因为除了上一个降尺度模块之外,之前的降尺度模块中的InvBlocks都会对空间分辨率更大的张量进行卷积运算,这大大增加了计算成本。为了进一步提高计算效率,我们建议修改架构,首先应用下采样模块(例如多个Haar变换或可学习的下采样),然后通过多个InvBlocks。这使得卷积运算可以应用在更小的分辨率上,这可以在相似数量的参数下大大降低FLOP和运行时间。
我们将这种架构下的IRN模型表示为IRNE。它与IRN的区别仅在于IRN堆叠了多个降尺度模块。
2.2 图像脱色-着色模型
图像脱色-着色是一种常见的任务,是双向退化-恢复问题的另一个实例,其中通道维度上的颜色信息丢失。我们的问题公式的核心思想与图1相同,它将特定于图像的丢失信息的分布转换为图像不可知的高斯分布。下面说明了为适应这项任务而对特定模型进行的一些调整
2.2.1 架构
其基本架构如图2所示。不同于图像重缩放时对低频和高频内容进行分割,我们需要对灰度信息和颜色信息进行分割,在这里获取颜色信息分布的同时生成灰度图像。因此,我们需要将下采样模块替换为灰化模块。我们直接利用图像的YCbCr颜色空间表示来分割通道中的信息。然后这两个信息分支(即Y和CbCr)通过前面介绍的InvBlock。我们将这个模型表示为IRNcolor。
2.3 图像重缩放与压缩的结合
我们的可逆框架将退化和恢复联合建模为可逆双射变换。在实际应用中,退化操作的某些部分并不总是可以与恢复相适应,例如为了方便。例如,广泛使用的图像压缩遵循通用标准,PNG和JPEG等格式是最常用的格式,在大多数数字设备中都得到了完善的支持。因此,我们建议将我们的可逆框架与现有退化方法的恢复相结合,以获得更广泛的应用。
具体来说,我们考虑图像重缩放和压缩结合的实例化,这也是高分辨率图像更高压缩率的常用方法,因为直接图像压缩方法在极大的压缩率下表现不佳。在这项工作中,我们演示了IRN和无损以及有损压缩方法之间的结合,以获得更好的有损压缩性能。
2.3.1 方法
在无损图像压缩方法中,LR图像可以在编码时不丢失信息,因此IRN可以直接与LR图像结合,即直接压缩IRN生成的下采样LR图像。
对于现有的有损图像压缩方法,在编码过程中不可避免地会有信息丢失,即有损压缩导致额外的退化。因此,直接将IRN与它们结合,例如先压缩IRN的LR图像,然后将压缩后的图像直接传递给IRN,可能违背了所提出的可逆框架中建模丢失信息的原则。为了获得良好的性能,需要对这种退化进行额外的恢复。
为了缓解这个问题,我们建议利用一个额外的模块,通过有损压缩方法部分恢复丢失的信息。具体来说,IRN的缩小图像将首先通过有损压缩方法(例如JPEG)进行压缩,压缩后的图像将通过压缩恢复模块(CRM)传递给IRN。采用CRM作为神经网络模型,输入是经过压缩后的退化LR图像,输出是经过有损压缩后退化恢复的LR图像。该模块被训练恢复给定压缩方法的丢失信息,这与考虑单向恢复任务的许多方法类似。我们将在下一节中详细阐述该架构并评估压缩性能。IRN和CRM的结合是我们提出的可逆框架和从现有退化方法恢复的结合的实例化模型。
三、实现细节
我们的实验包括三个部分:图像重缩放,图像的脱色-着色,以及图像缩放与压缩的结合。对于所有任务的训练,我们使用广泛使用的DIV2K图像恢复数据集来训练我们的模型。它包含800个高质量的2K分辨率的训练图像和100个验证图像。此外,对于前两项任务,我们在4个额外的标准数据集上评估我们的模型,即Set5,Set14,BSD100和Urban100 ;对于第三个任务,我们还在广泛使用的Kodak数据集上评估我们的模型。对于图像重缩放,按照(Lim等,2017)中的设置,我们定量评估了在YCbCr (Y,Cb,Cr)颜色空间中表示的图像Y通道上的峰值信噪比(PSNR)和SSIM。我们还将LPIPS、PI和FID作为感知评价的定量指标。对于另外两个任务,我们在RGB三通道颜色空间上评估PSNR和SSIM。
对于图像的重缩放,我们分别用1、2和3个降尺度模块在2×、4×和8×下采样下训练和比较我们的IRN模型。每个降尺度模块有8个InvBlocks,将原始图像下采样2倍。InvBlocks中的变换函数φ(·)、η(·)、ρ(·)由一个密集连接的卷积块进行参数化,Wang et al(2018)将其称为密集块。对于IRNLD模型在3×尺度下的实验,我们使用了一个降尺度模块和12个InvBlocks。对于IRNE模型在4×尺度下的实验,我们使用一个带16个Invblocks的降尺度模块(先下采样)。我们使用Adam优化器, β1 = 0.9, β2 = 0.999来训练我们的模型。mini-batch大小设置为16。输入的HR图像被随机裁剪为144 × 144,并通过应用随机的水平和垂直翻转来增强。在预训练阶段,迭代总次数为500K,学习率初始化为2 × 10−4,在[100k,200k,300k,400k]次小批量更新时减半。式(12)中的超参数设为λ1 = 1, λ2 = s2, λ3 = 1,其中s为标度。在预训练之后,我们对我们的模型进行微调,以进行另一个200K迭代,如3.3.3节所述。学习率初始化为1 × 10−4,并在[50k,100k]次迭代中减半。在式(13)中设置λ1 = 0.01, λ2 = s2, λ3 = 1, λ4 = 0.01,预训练鉴别器5000次迭代。鉴别器类似于Ledig等(2017),它包含8个具有3×3内核的卷积层,其数量每两层从64增加到512,增加2倍,以及两个具有100个隐藏单元的密集层。
对于图像的脱色-着色,灰色模块有8个InvBlocks。超参数设置为λ1 = 1, λ2 = 9, λ3 = 1。其他优化器和迭代设置与图像缩放相同。
为了与图像压缩相结合,我们利用在图像重缩放任务中训练的IRN2×模型,并通过在训练中缩放时对生成的LR图像添加随机噪声,进一步在缩放任务中对其进行100K迭代微调,以使模型对由于压缩和恢复而对LR图像可能的变化更具鲁棒性。Kodak的模型是额外微调的2.5K迭代Kodak。我们训练一个压缩恢复模块(CRM)为每个压缩比的JPEG。CRM包含中提出的残差密集块(RRDB)中的8个残差,并通过对重构LR图像和压缩前LR图像的L2损失进行训练。优化器和迭代设置与IRN相同。
四、代码运行
4.1 官方README文档
依赖关系和安装
(1)Python 3(推荐使用Anaconda)
(2)PyTorch >= 1.0
(3)NVIDIA GPU + CUDA
(4)Python包:pip install numpy opencv-python lmdb pyyaml
(5)TensorBoard:
PyTorch >= 1.1: pip install tb-nightly future
PyTorch == 1.0: pip install tensorboardX
图像缩放训练
首先在options/train/中设置一个配置文件,然后按如下方式运行:
python train.py -opt options/train/train_IRN_x4.yml
图像缩放测试
首先在options/test/中设置一个配置文件,然后按如下方式运行:
python test.py -opt options/test/test_IRN_x4.yml
图像脱色-着色训练
首先在options/train/中设置一个配置文件,然后按如下方式运行:
python train_IRN-Color.py -opt options/train/train_IRN_color.yml
图像脱色-着色测试
首先在options/test/中设置一个配置文件,然后按如下方式运行:
python test_IRN-Color.py -opt options/test/test_IRN_color.yml
结合图像压缩训练
首先在options/train/中设置一个配置文件,然后按如下方式运行:
python train_IRN-Compression.py -opt options/train/train_IRN-Compression_x2_q90.yml
结合图像压缩测试
首先在options/test/中设置一个配置文件,然后按如下方式运行:
python test_IRN-Compression.py -opt options/test/test_IRN-Compression_x2_q90.yml
代码框架
代码框架遵循BasicSR。它主要由配置、数据、模型和网络四部分组成。
让我们使用train命令python train.py -opt options/train/train_IRN_x4为例。在此命令之后将执行一系列操作。
(1)train.py被调用。
(2)读取选项/train/train_IRN_x4.yml中的配置,包括数据加载器的配置,网络,丢失,培训策略等。配置文件由options/options.py处理。
(3)创建训练和验证数据加载器。数据加载器是根据不同的数据模式在data/__init__.py中构造的。
(4)创建模型(根据不同的模型类型在models/__init__.py中构造)。
(5)开始训练模型。其他操作如日志记录、保存中间模型、验证、更新学习率等也在训练过程中完成。
4.2 实际运行代码
4.2.1 配置过程
(1)将codes目录设置为源代码根目录(pycharm)
(2)安装lmdb包:conda install -c conda-forge python-lmdb
(3)安装opencv包:conda install -c https://conda.anaconda.org/menpoopencv
(4)安装yaml包:conda install pyyaml
(5)安装tensorboardX包:conda install tensorboardX(pycharm)
(6)安装seaborn包:conda install seaborn(pycharm)
4.2.2 运行过程
(1)切换目录:cd codes
(2)执行命令:python test.py -opt options/test/test_IRN_x4.yml(注意注释其余test)
五、汇报PPT
论文地址:Invertible Rescaling Network and Its Extensions
开源代码:IRN+