DIDN(Deep Iterative Down-Up CNN for Image Denoising)

在网络中使用放大和缩小的特征图在低层次视觉任务中被广泛研究。作者提出了用于图像去噪的深度迭代缩放卷积网络(DIDN)，在网络中重复的减小和增大特征图的分辨率。其基本结构受用于语义分割的UNet启发，作者修改了缩放层使得其可用在图像去噪任务中。
作者指出了现有方法在结构和训练过程上的局限性。首先，图像去噪任务中需要输出与输入图像拥有相同的分辨率，这意味着网络中的深度特征需要与输出有相同的分辨率，这消耗了大量的GPU内存和训练时间。这些GPU和训练时间成本解释了现有的简单CNN去噪结构在其深度，参数量和感受野方面的限制。使用改变特征图分辨率的分层网络结构可以实现性能改进，因为对于相同的GPU内存成本，感受野可以更大。
MWCNN通过将小波变换和逆小波变换方法与U-Net结构相结合，搭建了一个图像重建网络，实现了计算复杂度和感受野间的良好平衡。然而，由于小波变换由具有特定权重的卷积组成，和子采样过程，它被认为是卷积层的一个特例; 与使用可训练卷积层可以实现的性能相比，使用这种方法可能会限制性能。
为了处理这种局限，作者将缩小放到压缩过程，使用步长为2的卷积层实现，将放大放到扩张过程，使用sub-pix实现。如果在放大之前提供足够数量的特征，则可以减少在放大过程期间的丢失。作者还使用了权重平均方法（Snapshot ensembles: Train 1, get m for free）去减小由权重选择造成的偏差，在不增加参数量的条件下提高了模型的性能。
该工作有以下贡献：

• 新结构：使用迭代的压缩和扩张特征，感受野很大
• 应用了权重平均法

算法细节

DIDN结构

DIDN.png

图示说明.png

DIDN整体架构如上图所示，灰色块代表特征图，特征图由4个不同分辨率等级组成。DIDN有4个部分：特征提取，DUB(down-up block, 缩放块)，重建，增强。

• 初始特征提取(Initial feature extraction)：DIDN首先使用3x3卷积从大小为HxW的输入图像提取N(论文中是128)通道特征，然后用3x3、步长为2的卷积层提取出W/2 x H/2 x 2N的特征。
• DUB(Down-Up Block)：提取的特征经过多个DUB(论文中是4个)的迭代缩放。具体见下节。
• 重建(Reconstruction)：受MemNet启发，作者在最后一个DUB后接了相同的重建块去获取所有的局部输出的优点。所有DUB的输出构成了重建块的输入，所有重建块的输出被concat到一起用于下一步增强。重建块由9个Conv+PReLU组成。更具体地说，有四个由Conv+PReLU+Conv+PReLU组成残差块组成，结构最后有多的一个Conv。
• 增强(Enhancement)：最后，通过1x1卷积，输出的特征图的通道数被重组减小，并使用上采样去产生最后的去噪图像。

DUB结构

DUB.png

However, these upscaling layers contain interpolation or padding processes which can include degradation in the feature maps.
然而这些上采样方法都包含插值法和填充过程，会带有特征图的降解。

由于图像去噪是低层次视觉任务，在其中增加像素级别的精确性是很重要的，在DIDN中，采取了subpix卷积层所谓上采样操作。subpix不需要插值法和填充过程，而允许网络从低像素传播细节信息到高像素，这是一个在图像去噪上的优点。
这里也有其他的降采样方法，比如最大池化和sub sample，但是这里采用可训练的卷积层去提高性能。

其他

论文中实验部分提到，初始通道数设定的是128，DUB个数为6，参数量约为190M。
该论文除了结构之外的部分，和DHDN有极大的相似性，毕竟是出自相同的作者。所以对于相关工作、训练、实验结果等细节，可以直接参考上一篇文章DHDN论文总结，在此不再赘述。
下面展示自己实现的DIDN结构并给出具体代码。

DIDN_keras.png

代码实现

结构分析

在DHDN中，总体结构较为清晰，基本模块有较高内聚性。
在DIDN中，DUB作为基本块，其中还运用了缩放。所以应该先把缩放块写好，进而组成DUB，然后把重建和增强模块写出，最后搭建主干网络。

引包和计数器

from keras.layers import Input,PReLU,Conv2D,Add,Concatenate
from keras.layers import MaxPooling2D,UpSampling2D
from keras.models import Model
from keras import backend as K
import keras 
import tensorflow as tf

def init_name_counter():
    name_counter = {}
    name_counter['DUB'] = 0
    name_counter['down'] = 0
    name_counter['up'] = 0
    name_counter['Conv2d_PReLU'] = 0
    name_counter['reconstruction'] = 0
    name_counter['enhancement'] = 0

    return name_counter
name_counter = init_name_counter()

Conv+PReLU

DIDN和DHDN不同，DHDN的结构有高度对称性，所以在代码实现时，基本结构多用循环代替。而DIDN基本结构不能用循环来写，所以会出现很多Conv+PReLU的代码，造成代码的冗余，所以需要将Conv+PReLU提出来。

def Conv2d_PReLU(filter):
    def wrapper(inputs):
        with tf.name_scope('Conv2d_PReLU'+str(name_counter['Conv2d_PReLU'])):
            name_counter['Conv2d_PReLU']+=1
            x = Conv2D(filter,3,padding='same')(inputs)
            x = PReLU(shared_axes=[1, 2])(x)
        return x
    return wrapper

DUB

def DUB(filter):
    def wrapper(inputs):
        with tf.name_scope('DUB'+str(name_counter['DUB'])):
            name_counter['DUB']+=1
            level_outputs = []
            level_outputs.append(inputs)
            
            x = Conv2d_PReLU(filter)(inputs)
            x = Conv2d_PReLU(filter)(x)
            x = Add()([inputs,x])
            level_outputs.append(x)
            
            inputs = downsampling_block(filter*2)(x)
            x = Add()([inputs,Conv2d_PReLU(filter*2)(inputs)])
            level_outputs.append(x)

            inputs = downsampling_block(filter*4)(x)
            x = Add()([inputs,Conv2d_PReLU(filter*4)(inputs)])
            x = Conv2D(filter*8,1,padding='same')(x)

            inputs = Concatenate()([level_outputs[-1],upsampling_block(filter*2)(x)])
            x = Conv2D(filter*2,1,padding='same')(inputs)
            x = Add()([x,Conv2d_PReLU(filter*2)(x)])
            x = Conv2D(filter*4,1,padding='same')(x)

            inputs = Concatenate()([level_outputs[-2],upsampling_block(filter)(x)])
            inputs = Conv2D(filter,1,padding='same')(inputs)
            x = Conv2d_PReLU(filter)(inputs)
            x = Conv2d_PReLU(filter)(x)
            x = Add()([inputs,x])
            x = Conv2d_PReLU(filter)(x)
            x = Add()([level_outputs[-3],x])
        return x
    return wrapper

DUB.png

reconstruction

def reconstruction(filter):
    def wrapper(inputs):
        with tf.name_scope('reconstruction'+str(name_counter['reconstruction'])):
            name_counter['reconstruction']+=1
            for _ in range(4):
                x = Conv2d_PReLU(filter)(inputs)
                x = Conv2d_PReLU(filter)(x)
                inputs = Add()([inputs,x])
            x = Conv2d_PReLU(filter)(inputs)
        return x
    return wrapper

reconstruction.png

enhancement

def enhancement(filter):
    def wrapper(inputs):
        with tf.name_scope('enhancement'+str(name_counter['enhancement'])):
            name_counter['enhancement']+=1
            x = Conv2D(filter,1,padding='same')(inputs)
            x = Add()([x,Conv2d_PReLU(filter)(x)])
            x = upsampling_block(filter)(x)
        return x
    return wrapper

enhancement.png

DIDN

def DIDN():
    input_channel = 3
    input_shape = (64,64,input_channel)
    init_filter = 128
    DUB_number = 6
    
    inputs = Input(shape=input_shape)
    x = Conv2d_PReLU(init_filter)(inputs)
    x = downsampling_block(init_filter*2)(x)

    # DUBs
    DUB_outputs = []
    for _ in range(DUB_number):
        # x is every DUB's output to reconstruction
        x = DUB(init_filter*2)(x)
        DUB_outputs.append(x)
    
    # I concatenate all DUB_outputs before reconstruction, 
    # which is different from the paper's description 
    # as concatenate all reconstruction's outputs 
    # I don't know how to implement the same architecture as the paper.
    x = Concatenate()(DUB_outputs)
    x = reconstruction(init_filter*2*DUB_number)(x)
    x = enhancement(init_filter*2)(x)
    x = Conv2d_PReLU(input_channel)(x)
    outputs = Add()([inputs,x])

    model = Model(inputs=inputs,outputs=outputs)
    return model