pytorch学习以FSRCNN为例

 

目录

 

 前言

一、FSRCNN网络结构

二.网络模型

三.数据集构建

四.训练

        4.1训练代码

        4.2可视化

五.验证

总结


 

 前言

本人刚刚开始学习pytorch,又刚好学习了FSRCNN,想着已学习的代码知识来练习写一个基于FSRCNN的图像超分辨,只记录代码编写过程,没有优化,指标也达不到论文效果,单纯记录学习过程.https://arxiv.org/pdf/1608.00367.pdf

一、FSRCNN网络结构

  FSRCNN是汤晓鸥团队为了改进SRCNN提出的,FSRCNN是采用转置卷积代替了SRCNN中的插值放大,简化了运算量,并且将SRCNN中的非线性映射步骤分解为了降维、映射和扩展三步,采用更深的卷积层,并且采用更小的卷积核,不仅降低了计算量,还增加了图像重建的质量。用以下图3.2 FSRCNN网络结构图。在图中,将卷积层代表为eq?Conv%28f%2Cn%2Cc%29,转置卷积代表eq?Deconv%28f%2Cn%2Cc%29,其中eq?f%2Cn%2Cc分别代表卷积核的大小,个数和通道数。  由图可知,FSRCNN整个网络分为5个步骤:特征提取、降维、非线性映射、扩展和转置卷积。下面对这五个步骤进行详细介绍:41bda7e79ab64165906e1f8d55000856.png

 

(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的完整模型如下:

 54d707ce03484417a8fc2fd3896ad108.png

二.网络模型

代码如下:手动敲,未优化代码结构
选取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:下图是训练过程

898d16721b354efda40c2f8a2a85b8a7.png

        4.2可视化

 

训练过程中加入了tensorboard:下图是tensorboard的可视化:

训练集损失下降

5dcbe1000e034af3b53ae32e1c9ce58b.png

 

 测试集平均损失下降

38a0f42b973b46f981b472bc01e1fd48.png

 

 平均PSNR上升

329bd85f36e847faa242d516b6722ee7.png

 

五.验证

验证代码如下:

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上的验证结果:

2967fa9b8e6f41e7b6c3257e3b4b51c3.png

 

 

总结

本文只是用来记录学习过程,简单复现FSRCNN模型但是并没有进行调优,未达到论文效果,

经过分析主要有以下原因:(1)首先采用Resize进行下采样默认好像不是双三次插值

(2)BSD300训练集中只有200张图片,数据集过少,导致泛化能力不好

(3)训练次数较少,模型通过tensorboard没完全收敛(不过就是收敛了,PScompression

(4)代码写的有些烂,问题有点大,但是不想改了。。。

 

 

  • 2
    点赞
  • 18
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

theshy123333

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值