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链接
