CVPR19-Deep Stacked Hierarchical

CVPR19-Deep Stacked Hierarchical Multi-patch Network for Image Deblurring论文复现

该工作主要关注于利用深度网络来实现图片去模糊，这里我们针对GoPro数据集进行论文的复现。

文章给出了一种新的模型架构，来学习不同层次上的特征，并实现去模糊的效果。

首先这里我们给出整体模型的架构

模型架构

如上图所示，整个模型由4个编码解码器构成，自底向上进行传播。

可以看到，从最下面的输入开始，我们将模糊的图片进行输入，会将图片分成8个区域，每个区域过编码器4，得到8个中间的特征表示，将8个中间的特征表示进行两两连接得到4个特征表示，并输入到解码器4，进而得到4个输出。

在输入到编码器3之前，会将图片分4块，然后将之前解码器4的输出加到图片的4块区域上，通过编码器3得到4个中间特征表示，这里我们将4个中间特征表示和先前两两连接的特征表示进行相加， 相加之后两两连接得到2个特征表示，并输入到解码器3。

依此进行反复操作，直到最上层，将最终的输出和对应的清晰的图片进行均方误差MSE的计算，然后反向传播进行模型的训练。

这里给出对应的pytorch代码

import torch.nn as nn
import torch
from decoder import Decoder
from encoder import Encoder

class DMPHNModel(nn.Module):
    def __init__(self, level=4, device='cuda'):
        super(DMPHNModel, self).__init__()
        self.encoder1 = Encoder().to(device)
        self.decoder1 = Decoder().to(device)
        self.encoder2 = Encoder().to(device)
        self.decoder2 = Decoder().to(device)
        self.encoder3 = Encoder().to(device)
        self.decoder3 = Decoder().to(device)
        self.encoder4 = Encoder().to(device)
        self.decoder4 = Decoder().to(device)
        self.level = level

    def forward(self, x):
        # x structure (B, C, H, W)
        # from bottom to top
        tmp_out = []
        tmp_feature = []
        for i in range(self.level):
            currentlevel = self.level - i - 1  # 3,2,1,0
            # For level 4(i.e. i = 3), we need to divide the picture into 2^i parts without any overlaps
            num_parts = 2 ** currentlevel
            rs = []
            if currentlevel == 3:
                rs = self.divide(x, 2, 4)
                for j in range(num_parts):
                    tmp_feature.append(self.encoder4(rs[j]))  # each feature is [B, C, H, W]
                # combine the output
                tmp_feature = self.combine(tmp_feature, comb_dim=3)
                for j in range(int(num_parts/2)):
                    tmp_out.append(self.decoder4(tmp_feature[j]))
            elif currentlevel == 2:
                rs = self.divide(x, 2, 2)
                for j in range(len(rs)):
                    rs[j] = rs[j] + tmp_out[j]
                    tmp_feature[j] = tmp_feature[j] + self.encoder3(rs[j])
                tmp_feature = self.combine(tmp_feature, comb_dim=2)
                tmp_out = []
                for j in range(int(num_parts/2)):
                    tmp_out.append(self.decoder3(tmp_feature[j]))
            elif currentlevel == 1:
                rs = self.divide(x, 1, 2)
                for j in range(len(rs)):
                    rs[j] = rs[j] + tmp_out[j]
                    tmp_feature[j] = tmp_feature[j] + self.encoder2(rs[j])
                tmp_feature = self.combine(tmp_feature, comb_dim=3)
                tmp_out = []
                for j in range(int(num_parts/2)):
                    tmp_out.append(self.decoder2(tmp_feature[j]))
            else:
                x += tmp_out[0]
                x = self.decoder1(self.encoder1(x)+tmp_feature[0])
        return x
    
    def combine(self, x, comb_dim=2):
        """[将数组逐两个元素进行合并并且返回]

        Args:
            x ([tensor array]): [输出的tensor数组]
            comb_dim (int, optional): [合并的维度，从高度合并则是2，宽度合并则是3]. Defaults to 2.

        Returns:
            [tensor array]: [合并后的数组，长度变为一半]
        """        
        rs = []
        for i in range(int(len(x)/2)):
            rs.append(torch.cat((x[2*i], x[2*i+1]), dim=comb_dim))
        return rs

    def divide(self, x, h_parts_num, w_parts_num):
        """ 该函数将BxHxWxC的输入进行切分, 本质上是对每一张图片进行分块
            这里直接针对多维数组进行操作

        Args:
            x (Torch Tensor): input torch tensor (e.g. [Batchsize, Channels, Heights, Width])
            h_parts_num (int): The number of divided parts on heights
            w_parts_num (int): The number of divided parts on width

        Returns:
            [A list]: h_parts_num x w_parts_num 's tensor list, each one has [B, Channels, H/h_parts_num, W/w_parts_num] structure
        """                
        rs = []
        for i in range(h_parts_num):
            tmp = x.chunk(h_parts_num, dim=2)[i]
            for j in range(w_parts_num):
                rs.append(tmp.chunk(w_parts_num,dim=3)[j])
        return rs

上述代码是整个模型的输入流程，我们还需要对其中的编码解码器的结构进行实现
这里给出论文中的结构描述

编码器和解码器

这里论文给出的图片有些错误，解码器的最后一层应该是[32,3,3,1]，不然无法输出3个通道的图片。

灰色块的ReLU激活函数，块和块之间的有向连接是残差连接，是先做卷积再做加法

直接给出对应的编码器和解码器的代码

import torch.nn as nn
import torch.nn.functional as F

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1, stride=1)
        self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1)
        self.conv3 = nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1)   
        self.conv4 = nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1)
        self.conv5 = nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1)
        self.conv6 = nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=2)
        self.conv7 = nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1)
        self.conv8 = nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1)
        self.conv9 = nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1)
        self.conv10 = nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1)
        self.conv11 = nn.Conv2d(64, 128, kernel_size=3, padding=1, stride=2)
        self.conv12 = nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1)
        self.conv13 = nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1) 
        self.conv14 = nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1) 
        self.conv15 = nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1)           
        
    def forward(self, x):
        tmp = self.conv1(x)
        x1 = F.relu(self.conv2(tmp))
        x1 = self.conv3(x1)
        tmp = x1 + tmp  # residual link
        x1 = F.relu(self.conv4(tmp))
        x1 = self.conv5(x1)
        x1 = x1 + tmp  # residual link
        tmp = self.conv6(x1)
        x1 = F.relu(self.conv7(tmp))
        x1 = self.conv8(x1)
        tmp = x1 + tmp  # residual link
        x1 = F.relu(self.conv9(tmp))
        x1 = self.conv10(x1)
        x1 = x1 + tmp  # residual link
        tmp = self.conv11(x1)
        x1 = F.relu(self.conv12(tmp))
        x1 = self.conv13(x1)
        tmp = x1 + tmp  # residual link
        x1 = F.relu(self.conv14(tmp))
        x1 = self.conv15(x1)
        x1 = x1 + tmp  # residual link
        return x1

解码器的代码为：

import torch.nn as nn
import torch.nn.functional as F

class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.conv1 = nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1)
        self.conv2 = nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1)
        self.conv3 = nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1)
        self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=1)
        self.deconv1 = nn.ConvTranspose2d(128, 64, kernel_size=4, padding=1, stride=2)
        self.conv5 = nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1)
        self.conv6 = nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1)
        self.conv7 = nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1)
        self.conv8 = nn.Conv2d(64, 64, kernel_size=3, padding=1, stride=1)
        self.deconv2 = nn.ConvTranspose2d(64, 32, kernel_size=4, padding=1, stride=2)
        self.conv9 = nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1)
        self.conv10 = nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1)
        self.conv11 = nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1)
        self.conv12 = nn.Conv2d(32, 32, kernel_size=3, padding=1, stride=1)
        self.conv13 = nn.Conv2d(32, 3, kernel_size=3, padding=1, stride=1)
    
    def forward(self, x):
        tmp = x
        x1 = F.relu(self.conv1(tmp))
        x1 = self.conv2(x1)
        tmp = x1 + tmp  # residual link
        x1 = F.relu(self.conv3(tmp))
        x1 = self.conv4(x1)
        x1 = x1 + tmp  # residual link
        tmp = self.deconv1(x1)
        x1 = F.relu(self.conv5(tmp))
        x1 = self.conv6(x1)
        tmp = x1 + tmp  # residual link
        x1 = F.relu(self.conv7(tmp))
        x1 = self.conv8(x1)
        x1 = x1 + tmp  # residual link
        tmp = self.deconv2(x1)
        x1 = F.relu(self.conv9(tmp))        
        x1 = self.conv10(x1)
        tmp = x1 + tmp  # residual link
        x1 = F.relu(self.conv11(tmp))
        x1 = self.conv12(x1)
        x1 = x1 + tmp  # residual link
        return self.conv13(x1)

这里我训练了1500个epoch，lr设置为1e-4，batch_size=6。训练完的效果如下图所示

Input	Output
input	在这里插入图片描述

完整的训练代码我已经放到了github上，欢迎大家star和issue，链接如下：
github链接