目录
前言
本人刚刚开始学习pytorch,又刚好学习了FSRCNN,想着已学习的代码知识来练习写一个基于FSRCNN的图像超分辨,只记录代码编写过程,没有优化,指标也达不到论文效果,单纯记录学习过程.https://arxiv.org/pdf/1608.00367.pdf
一、FSRCNN网络结构
FSRCNN是汤晓鸥团队为了改进SRCNN提出的,FSRCNN是采用转置卷积代替了SRCNN中的插值放大,简化了运算量,并且将SRCNN中的非线性映射步骤分解为了降维、映射和扩展三步,采用更深的卷积层,并且采用更小的卷积核,不仅降低了计算量,还增加了图像重建的质量。用以下图3.2 FSRCNN网络结构图。在图中,将卷积层代表为,转置卷积代表,其中分别代表卷积核的大小,个数和通道数。 由图可知,FSRCNN整个网络分为5个步骤:特征提取、降维、非线性映射、扩展和转置卷积。下面对这五个步骤进行详细介绍:
(1)特征提取,这一过程就是提取图像的特征,在SRCNN中是双三次插值采用大小的卷积核提取特征,而FSRCNN则是直接输入LR图像,采用大小(5×5)的卷积核进行卷积。由于插值的缘故,FSRCNN中采用较小的卷积核提取LR图像的特征与SRCNN中放大后提取的特征信息几乎相同。所以这样做减少了计算量并且,提取的特征几乎没有损失。
(2)降维,在特征提取的过程中一般会使用较多的卷积核来提取特征,上一层卷积核的数量是这一层输入的维度,所以特征提取后的维度过大直接进行非线性映射会让计算复杂度过高,所以采用较少个数的(1×1)大小卷积核来进行降维处理,目的就是降低计算开销。
(3)非线性映射,非线性映射对图像重建质量的影响很大,FSRCNN中使用多个映射层,采用大小适中的(3×3)大小的卷积核代替SRCNN中的(1×1)卷积核,通常情况下,映射层的每个卷积核的数量应该相同。
(4)扩展,扩展的操作与第二步的降维刚好相反。降维的目的是降低计算量,得到低维的特征图,但是如果直接用低维度的特征图进行图像重建,则会大大降低重建的质量,所以需要扩展升维,依然采用大小(1×1)的卷积核。
(5)转置卷积,转置卷积的作用就是实现上采样对图像进行放大,转置卷积不同于插值等一些上采样的方法,他是对输入特征进行上采样的卷积核,是可以自动学习的。通过转置卷积即可得到HR图像。
FSRCNN的模型中在每个卷积层后采用了的PreLU函数代替了ReLU函数,PreLU函数防止了出现梯度为0的情况,事实证明,使用PreLU函数的网络稳定性更高,在PreLU函数中的是一个可以学习的参数。总结以上五个步骤,一个FSRCNN的完整模型如下:
二.网络模型
代码如下:手动敲,未优化代码结构
选取d=64 s=16 m=4放大倍数为4
import torch
from torch import nn
from torch.nn import Conv2d, PReLU, ConvTranspose2d
"""
FSRCNN网络结构
"""
class YJL(nn.Module):
def __init__(self):
super(YJL, self).__init__()
self.conv1 = Conv2d(3 ,64 ,kernel_size=(5,5),padding=2)
self.PReLU1 = PReLU()
self.conv2 = Conv2d(64, 16, kernel_size=(1,1),padding=0)
self.PReLU2 = PReLU()
self.conv3 = Conv2d(16, 16, kernel_size=(3,3),padding=1)
self.PReLU3 = PReLU()
self.conv4 = Conv2d(16, 16, kernel_size=(3,3),padding=1)
self.PReLU4 = PReLU()
self.conv5 = Conv2d(16, 16, kernel_size=(3,3),padding=1)
self.PReLU5 = PReLU()
self.conv6 = Conv2d(16, 16, kernel_size=(3,3),padding=1)
self.PReLU6 = PReLU()
self.conv7 = Conv2d(16, 64 ,kernel_size=(1,1))
self.PReLU7 = PReLU()
self.conv8 = ConvTranspose2d(64, 3, kernel_size=(9,9),stride=4 ,padding=4,output_padding=3)
self.PReLU8 = PReLU()
def forward(self, x):
x = self.conv1(x)
x = self.PReLU1(x)
x = self.conv2(x)
x = self.PReLU2(x)
x = self.conv3(x)
x = self.PReLU3(x)
x = self.conv4(x)
x = self.PReLU4(x)
x = self.conv5(x)
x = self.PReLU5(x)
x = self.conv6(x)
x = self.PReLU6(x)
x = self.conv7(x)
x = self.PReLU7(x)
x = self.conv8(x)
x = self.PReLU8(x)
return x
# scale_factor=4, num_channels=3, d=48, s=12, m=4).to(device)
"""
调试验证,输入40输出应为160大小
"""
if __name__ == '__main__':
yjl = YJL()
input = torch.ones(64,3 ,40 ,40)
output = yjl(input)
print(output.shape)
三.数据集构建
代码如下(示例):
import os
import torch
import torchvision.transforms
from PIL import Image
"""
创建数据集:
数据集创建整体思路:
将图片随机裁剪至固定大小,再通过Resize下采样得到低分辨率图像,返回这两种对于HR,LR图像
"""
class Mydataset:
def __init__(self, train_dir, patch_size): #训练集地址,和patchsize
self.train_dir = train_dir
self.train_path = os.listdir(self.train_dir) #获取文件地址(列表)
self.patch_size = torchvision.transforms.RandomCrop(patch_size) #采用transforms随机裁剪大小为patchsize
def __getitem__(self, idx):
img_name = self.train_path[idx] #获取各个图片索引
img_item_path = os.path.join(self.train_dir,img_name) #地址相加获取文件中各个图片的地址
img_PIL = Image.open(img_item_path)
if img_PIL.mode != 'RGB':
img_PIL = img_PIL.convert('RGB') #转换为RGB
img_tensor = torchvision.transforms.ToTensor()(img_PIL) #将PIL图片类型转换为tensor
img_corp = self.patch_size(img_tensor) #将其随机裁剪为patchsize大小
img_HR = img_corp
img_LR = torchvision.transforms.Resize(40)(img_HR) # 将其得到的图像通过Resize裁剪为40大小,其中我选取的Patchsize选取为160大小,对应4倍关系
return img_HR ,img_LR #返回HR,和LR图像 类型为tensor
def __len__(self):
return len(self.train_path) #返回数据集长度
"""
调试验证
"""
if __name__ == '__main__':
train_dir = "BSDS300/images/train"
patch_size = 160
dataset = Mydataset(train_dir,patch_size)
imghr,imglr = dataset[0]
print(imghr.shape)
print(imglr.shape)
该处使用的url网络请求的数据。
四.训练
4.1训练代码
代码如下:
import torch.nn from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from model import * from dataset import * if __name__ == '__main__': device = torch.device("cuda") #采用GPU训练 train_dir = "BSDS300/images/train" #训练集地址 test_dir = "BSDS300/images/test" #测试集地址(训练集数据集采用与训练集相同方式) """ dataloader:加载定义的数据集 """ patch_size = 160 dataset_train = Mydataset(train_dir,patch_size) dataloader_train = DataLoader(dataset_train,batch_size=64,shuffle=False) dataset_test = Mydataset(test_dir,160) dataloader_test = DataLoader(dataset_test,batch_size=64,shuffle=False) writer = SummaryWriter("logggg") #加入tensorboard可视化 """ 输出训练集和测试的图像数量 """ date_train_size = len(dataset_train) date_test_size = len(dataset_test) print("训练集长度为:{}".format(date_train_size)) print("测试集长度为:{}".format(date_test_size)) """ 创建神经网络 """ yuan = YJL() yuan = yuan.to(device) """ 定义损失函数 """ loss_func = torch.nn.MSELoss() loss_func = loss_func.to(device) """ 定义psnr """ def torchPSNR(tar_img, prd_img): imdff = torch.clamp(prd_img,0,1) - torch.clamp(tar_img,0,1) rmse = (imdff**2).mean().sqrt() ps = 20*torch.log10(1/rmse) return ps """ 定义优化器 """ learn_rate = 2e-5 optimizer = torch.optim.Adam(params=yuan.parameters(),lr=learn_rate) """ 开始训练 """ psnr_test = 0 total_step = 0 epoch = 1000#训练轮数 for i in range(epoch): print("-------第{}轮训练开始--------".format(i+1)) for data in dataloader_train: img_hr ,img_lr = data img_hr = img_hr.to(device) img_lr = img_lr.to(device) output = yuan(img_lr) loss = loss_func(output,img_hr) """ 优化器 """ optimizer.zero_grad() loss.backward() #反向传播 optimizer.step()#梯度下降 total_step += 1 writer.add_scalar("loss", loss ,total_step) #将loss可视化 #if total_step % 64 ==0: print("--第{}次训练的loss为:{}--".format(total_step,loss.item())) with torch.no_grad(): total_loss = 0 total_psnr = 0 for data in dataloader_test: img_hr, img_lr = data img_hr = img_hr.to(device) img_lr = img_lr.to(device) output = yuan(img_lr) loss = loss_func(output ,img_hr) """ 计算psnr """ psnr = torchPSNR(output,img_hr) total_loss += loss total_psnr += psnr print("测试集平均损失为:{}".format(total_loss/len(dataloader_test))) print("测试集平均psnr为:{}".format(total_psnr/len(dataloader_test))) writer.add_scalar("psnr",total_psnr/len(dataloader_test),i+1) #可视化 writer.add_scalar("ave_loss",total_loss/len(dataloader_test),i+1)#可视化 if total_psnr/date_test_size > psnr_test: torch.save(yuan,"yuan_FSRCNN_64_16_4.pth") print("模型已保存") #若PSNR比上一次提高则保存模型 psnr_test = total_psnr/date_test_size writer.close()
采用BDS300数据集训练和测试 代码中epoch=100,在我训练时改为了1000:下图是训练过程
4.2可视化
训练过程中加入了tensorboard:下图是tensorboard的可视化:
训练集损失下降
测试集平均损失下降
平均PSNR上升
五.验证
验证代码如下:
import os import torch import torchvision.transforms from PIL import Image device = torch.device("cuda") """ 采用Set5数据集测试 """ LR_file_path ="LRbicx4" #x4为GTmod12图片4倍数下采样后的图片 HR_file_path = "GTmod12" module = torch.load("yuan_64.pth")#加载保存的文件 print(module) img_LR_name = os.listdir(LR_file_path) img_HR_name = os.listdir(HR_file_path) """ 定义PSNR,也可以调用train.py的公式 """ def torchPSNR(tar_img, prd_img): imdff = torch.clamp(prd_img, 0, 1) - torch.clamp(tar_img, 0, 1) rmse = (imdff ** 2).mean().sqrt() ps = 20 * torch.log10(1 / rmse) return ps """ 获取长度 """ LR_size = len(img_LR_name) HR_size = len(img_HR_name) if LR_size != HR_size: print("文件图片数量不匹配") """ 通过For循环获取每一个文件 """ total_test_psnr = 0 for i in range(LR_size): img_HR_path = os.path.join(HR_file_path,img_HR_name[i]) img_LR_path = os.path.join(LR_file_path,img_LR_name[i]) img_LR_PIL = Image.open(img_LR_path) img_LR_tensor = torchvision.transforms.ToTensor()(img_LR_PIL).unsqueeze(dim=0).to(device) img_HR_PIL = Image.open(img_HR_path) img_HR_tensor = torchvision.transforms.ToTensor()(img_HR_PIL).unsqueeze(dim=0).to(device) img_SR = module(img_LR_tensor) psnr = torchPSNR(img_SR,img_HR_tensor) total_test_psnr +=psnr print("Set5的平均PSNR的值为:{}".format(total_test_psnr/LR_size))
在Set5上的验证结果:
总结
本文只是用来记录学习过程,简单复现FSRCNN模型但是并没有进行调优,未达到论文效果,
经过分析主要有以下原因:(1)首先采用Resize进行下采样默认好像不是双三次插值
(2)BSD300训练集中只有200张图片,数据集过少,导致泛化能力不好
(3)训练次数较少,模型通过tensorboard没完全收敛(不过就是收敛了,PScompression
(4)代码写的有些烂,问题有点大,但是不想改了。。。