IR_DIDN论文总结+结构实现

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

在DUB中,压缩和扩张路径由2个降采样和上采样完成。一个3x3、步长为2的卷积层和sub-pix层用作降采样和上采样。
降采样过程中特征图的边长减半,通道数倍增。上采样中,由于特征图通道减小到四分之一,所以在上采样之前先使用1x1的卷积层增加通道数以维持信息的密度。在UNet中,相同层级的特征被重用,这里开头和结尾的特征被连接。
还有其他的上采样方法存在,比如用反卷积或者Conv+upsampling。

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
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容