深度学习之超分辨率算法——VDSR

• 比较于之前的FSRCNN来说，VDSR我认为主要引入了以下优秀特点

• 首先是卷积层数的上的增加，卷积层数直接代表着模型提取特征的能力强弱
• 小卷积核的进一步引入，利用卷积核3x3堆叠层数，模型一共20层。
• 引入残差网络

• 缺点：

• 原文依然采用的MSE损失，单纯比较像素之间的差异。图像相对比较平滑。
• 训练依然采用是SRCNN的训练方法，先上采样到高分辨率尺寸大小再进行训练。

两张图理解：

（输入尺寸等于输出尺寸的模型实现）

• 使用数据集:train.h5

model.py

import torch
import torch.nn as nn
from math import sqrt

class VDSR(nn.Module):

    def __init__(self):
        super(VDSR, self).__init__()
        # 残差网络
        self.residual_layer = self.make_layer(Conv_ReLU_Block, 18)

        # 输入
        self.input = nn.Conv2d(in_channels=1, out_channels=64,  kernel_size=(3,3), stride=(1,1), padding=(1,1), bias=False)
        self.output = nn.Conv2d(in_channels=64, out_channels=1,  kernel_size=(3,3), stride=(1,1), padding=(1,1), bias=False)
        self.relu = nn.ReLU(inplace=True)

        self.init_weights()

    def init_weights(self):
        # 模型初始化参数
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, sqrt(2. / n))

    def make_layer(self, block, num_of_layer):
        # 做18层网络呀.....
        layers = []
        for _ in range(num_of_layer):
            layers.append(block())
        return nn.Sequential(*layers)

    def forward(self, x):
        residual = x
        # 1
        out = self.relu(self.input(x))
        # 18
        out = self.residual_layer(out)
        # 1
        out = self.output(out)
        # 残差
        out = residual+out
        return out


class Conv_ReLU_Block(nn.Module):

    def __init__(self):
        super(Conv_ReLU_Block, self).__init__()
        self.sequential = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.sequential(x)


if __name__ == '__main__':
    x=  torch.randn(1,1,224,224)
    net = VDSR()
    print(net(x).shape)      

dataset.py （数据读取方式）

import torch.utils.data as data
import torch
import h5py
from PIL import Image
import numpy as np


class DatasetFromHdf5(data.Dataset):
    def __init__(self, file_path="./data/train.h5"):
        super(DatasetFromHdf5, self).__init__()
        hf = h5py.File(file_path)
        self.data = hf.get('data')
        self.target = hf.get('label')

    def __getitem__(self, index):
        image = torch.from_numpy(self.data[index,:,:,:]).float()
        label = torch.from_numpy(self.target[index,:,:,:]).float()

        return image,label

        
    def __len__(self):
        return self.data.shape[0]

if __name__ == '__main__':
    data = DatasetFromHdf5()
    print(len(data))
    image = data[0][0]
    label = data[0][1]
    print(image.numpy()[0].shape)
    # 显示图片
    image = Image.fromarray(image.numpy()[0]*255)
    label = Image.fromarray(label.numpy()[0] * 255)
    image.show()
    label.show()      

• train.py

import argparse, os
import torch
import random
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from vdsr import VDSR
from dataset import DatasetFromHdf5

def init():
    # Training settings
    parser = argparse.ArgumentParser(description="the VDSR of Pytorch")

    # batch_size 每次投入模型的图像数据数
    parser.add_argument("--batch_size", type=int, default=128, help="Training batch size")
    # 训练轮次
    parser.add_argument("--epochs", type=int, default=1000, help="Number of epochs to train for")
    # 学习率
    parser.add_argument("--lr", type=float, default=0.0001, help="Learning Rate. Default=0.1")
    # 动态学习率调整系数
    parser.add_argument("--step", type=int, default=10, help="Sets the learning rate to the initial LR decayed by momentum every n epochs, Default: n=10")
    # 使用cuda
    parser.add_argument("--cuda", action="store_true",default=True,help="Use cuda?")
    # 已训练权重
    parser.add_argument("--resume", default="", type=str, help="Path to checkpoint (default: none)")
    # 开始轮次
    parser.add_argument("--start-epoch", default=1, type=int, help="Manual epoch number (useful on restarts)")
    # 梯度裁剪系数
    parser.add_argument("--clip", type=float, default=0.4, help="Clipping Gradients. Default=0.4")
    # 单线程
    parser.add_argument("--num_workers", type=int, default=1, help="Number of threads for data loader to use, Default: 1")
    # 优化器动量
    parser.add_argument("--momentum", default=0.9, type=float, help="Momentum, Default: 0.9")
    # 正则化系数
    parser.add_argument("--weight-decay", "--wd", default=1e-4, type=float, help="Weight decay, Default: 1e-4")
    #  预训练
    parser.add_argument('--pretrained', default='', type=str, help='path to pretrained model (default: none)')
    # 默认GPU为0
    parser.add_argument("--gpus", default="0", type=str, help="gpu ids (default: 0)")
    return parser

def main():

    parser = init()
    # 获得所有参数
    opt = parser.parse_args()
    print(opt)
    # cuda设置gpu参数
    cuda = opt.cuda

    # gpu配置
    if cuda:
        print("=> use gpu id: '{}'".format(opt.gpus))
        os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
        if not torch.cuda.is_available():
                raise Exception("No GPU found or Wrong gpu id, please run without --cuda")
    # 随机种子参数
    opt.seed = random.randint(1, 10000)
    print("Random Seed: ", opt.seed)
    torch.manual_seed(opt.seed)

    if cuda:
        # 设置固定生成随机数的种子，使得每次运行该.py
        # 文件时生成的随机数相同
        torch.cuda.manual_seed(opt.seed)

    # 设置加速，优化运行效率
    cudnn.benchmark = True

    print("===> Loading datasets")
    train_set = DatasetFromHdf5()

    # training_data_loader = DataLoader(dataset=train_set,num_workers=opt.num_workers, batch_size=opt.batch_size, shuffle=True)
    training_data_loader = DataLoader(dataset=train_set, batch_size=opt.batch_size,
                                      shuffle=True)
    print("===> Building model")
    model = VDSR()


    criterion = nn.MSELoss()

    print("===> Setting GPU")
    if cuda:
        model = model.cuda()
        criterion = criterion.cuda()

    # optionally resume from a checkpoint
    if opt.resume:
        if os.path.isfile(opt.resume):
            print("=> loading checkpoint '{}'".format(opt.resume))
            checkpoint = torch.load(opt.resume)
            opt.start_epoch = checkpoint["epoch"] + 1
            model.load_state_dict(checkpoint["model"].state_dict())
        else:
            print("=> no checkpoint found at '{}'".format(opt.resume))

    # optionally copy weights from a checkpoint
    if opt.pretrained:

        if os.path.isfile(opt.pretrained):
            print("=> loading model '{}'".format(opt.pretrained))
            weights = torch.load(opt.pretrained)
            model.load_state_dict(weights['model'].state_dict())
        else:
            print("=> no model found at '{}'".format(opt.pretrained))  

    print("===> Setting Optimizer")

    optimizer = optim.SGD(model.parameters(), lr=opt.lr, momentum=opt.momentum, weight_decay=opt.weight_decay)

    print("===> Training")

    for epoch in range(opt.start_epoch, opt.epochs + 1):
        train(opt,training_data_loader, optimizer, model, criterion, epoch)
        save_checkpoint(model, epoch)

def adjust_learning_rate(opt, epoch):
    """Sets the learning rate to the initial LR decayed by 10 every 10 epochs"""
    lr = opt.lr * (0.1 ** (epoch //opt.step))
    return lr

def train(opt,training_data_loader, optimizer, model, criterion, epoch):
    # lr = adjust_learning_rate(opt, epoch-1)
    lr = opt.lr
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    print("Epoch = {}, lr = {}".format(epoch, optimizer.param_groups[0]["lr"]))

    model.train()

    for iteration, batch in enumerate(training_data_loader):
        input, target = torch.Tensor(batch[0]), torch.Tensor(batch[1])

        if opt.cuda:
            input = input.cuda()
            target = target.cuda()

        # print("meodel:",model(input).shape)
        # print(target.shape)
        loss = criterion(model(input), target)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(),opt.clip)
        optimizer.step()

        if iteration%100 == 0:
            print("===> Epoch[{}]({}/{}): Loss: {:.10f}".format(epoch, iteration, len(training_data_loader), loss.item()))

def save_checkpoint(model, epoch):
    model_out_path = "checkpoint/" + "model_epoch_{}.pth".format(epoch)
    state = {"epoch": epoch ,"model": model}
    if not os.path.exists("checkpoint/"):
        os.makedirs("checkpoint/")

    torch.save(state, model_out_path)

    print("Checkpoint saved to {}".format(model_out_path))

if __name__ == "__main__":
    main()      

eval.py

import argparse, os
import torch
from torch.autograd import Variable
import numpy as np
import time, math, glob
import scipy.io as sio

parser = argparse.ArgumentParser(description="PyTorch VDSR Eval")
parser.add_argument("--cuda", action="store_true", help="use cuda?")
parser.add_argument("--model", default="model/model_epoch_100.pth", type=str, help="model path")
parser.add_argument("--dataset", default="Set5", type=str, help="dataset name, Default: Set5")
parser.add_argument("--gpus", default="0", type=str, help="gpu ids (default: 0)")

def PSNR(pred, gt, shave_border=0):
    height, width = pred.shape[:2]
    pred = pred[shave_border:height - shave_border, shave_border:width - shave_border]
    gt = gt[shave_border:height - shave_border, shave_border:width - shave_border]
    imdff = pred - gt
    rmse = math.sqrt(np.mean(imdff ** 2))
    if rmse == 0:
        return 100
    return 20 * math.log10(255.0 / rmse)

opt = parser.parse_args()
cuda = opt.cuda

if cuda:
    print("=> use gpu id: '{}'".format(opt.gpus))
    os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
    if not torch.cuda.is_available():
            raise Exception("No GPU found or Wrong gpu id, please run without --cuda")

model = torch.load(opt.model, map_location=lambda storage, loc: storage)["model"]

scales = [2,3,4]

image_list = glob.glob(opt.dataset+"_mat/*.*") 

for scale in scales:
    avg_psnr_predicted = 0.0
    avg_psnr_bicubic = 0.0
    avg_elapsed_time = 0.0
    count = 0.0
    for image_name in image_list:
        if str(scale) in image_name:
            count += 1
            print("Processing ", image_name)
            # 加载模型
            im_gt_y = sio.loadmat(image_name)['im_gt_y']
            # 双线性插值
            im_b_y = sio.loadmat(image_name)['im_b_y']
                       
            im_gt_y = im_gt_y.astype(float)
            im_b_y = im_b_y.astype(float)

            psnr_bicubic = PSNR(im_gt_y, im_b_y,shave_border=scale)
            avg_psnr_bicubic += psnr_bicubic


            im_input = im_b_y/255.

            im_input =torch.Tensor(torch.from_numpy(im_input).float()).reshape(1, -1, im_input.shape[0], im_input.shape[1])

            if cuda:
                model = model.cuda()
                im_input = im_input.cuda()
            else:
                model = model.cpu()

            start_time = time.time()
            HR = model(im_input)
            elapsed_time = time.time() - start_time
            avg_elapsed_time += elapsed_time

            HR = HR.cpu()

            im_h_y = HR.data[0].numpy().astype(np.float32)

            im_h_y = im_h_y * 255.
            im_h_y[im_h_y < 0] = 0
            im_h_y[im_h_y > 255.] = 255.
            im_h_y = im_h_y[0,:,:]

            psnr_predicted = PSNR(im_gt_y, im_h_y,shave_border=scale)
            avg_psnr_predicted += psnr_predicted

    print("Scale=", scale)
    print("Dataset=", opt.dataset)
    print("PSNR_predicted=", avg_psnr_predicted/count)
    print("PSNR_bicubic=", avg_psnr_bicubic/count)
    print("It takes average {}s for processing".format(avg_elapsed_time/count))      

像素重采样的VDSR

• 原文VDSR训练时，直接将原图上采样后"高分辨率"图像加入到模型中进行计算。
• 使用训练集VOC2012，RGB图
• 训练集目录如下：
• data存放数据假设bat_size = 1（1,3，224,224）
• SRF_2下：target:（1,3，448,448）
• SRF_3下：target(1,3,2243,2243)



• dataset.py

class DatasetFromVoc(data.Dataset):
    def __init__(self, file_path="./train",scale=2):
        super(DatasetFromVoc, self).__init__()
        if scale==2:
            dir_path = os.path.join(file_path,"SRF_2")
        elif scale==3:
            dir_path = os.path.join(file_path,"SRF_3")
        else:
            dir_path = os.path.join(file_path,"SRF_4")

        self.data_path = os.path.join(dir_path,"data")
        self.target_path = os.path.join(dir_path,"target")
        self.dataset = []
        for img_name in os.listdir(self.data_path):
            img_path = os.path.join(self.data_path,img_name)
            img_target = os.path.join(self.target_path,img_name)
            self.dataset.append([img_path,img_target])


    def __getitem__(self, index):
        img_path,label_path = self.dataset[index]
        img_data = cv.imread(img_path)
        label_data = cv.imread(label_path)
        # print(img_data.shape)
        # print(label_data.shape)
        img_data = img_data.transpose([2,0,1])
        label_data = label_data.transpose([2,0,1])

        img_data = np.array(img_data,dtype=np.float32)/255.
        label_data = np.array(label_data,dtype=np.float32)/255.
        return img_data,label_data

    def __len__(self):

        return len(self.dataset)      

vsnr.py

模型中我采取ESPCN的输出方式，采用像素混洗的方法，最后输出的时候才进行上采样，节约计算量。

进一步改进：损失函数应该加入感受损失

class VDSR_ESPCN(nn.Module):

    def __init__(self,input_channel=1,scale=2):
        super(VDSR_ESPCN, self).__init__()
        # 残差网络
        self.residual_layer = self.make_layer(Conv_ReLU_Block, 18)

        # 输入
        self.input = nn.Conv2d(in_channels=input_channel, out_channels=64,  kernel_size=(3,3), stride=(1,1), padding=(1,1), bias=False)
        self.output = nn.Conv2d(in_channels=64, out_channels=input_channel,  kernel_size=(3,3), stride=(1,1), padding=(1,1), bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.last_part = nn.Sequential(
            nn.Conv2d(input_channel, input_channel* (scale ** 2), kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            nn.PixelShuffle(scale)
        )

        self.init_weights()

    def init_weights(self):
        # 模型初始化参数
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, sqrt(2. / n))

    def make_layer(self, block, num_of_layer):
        # 做18层网络呀.....
        layers = []
        for _ in range(num_of_layer):
            layers.append(block())
        return nn.Sequential(*layers)

    def forward(self, x):
        residual = x
        # 1
        out = self.relu(self.input(x))
        # 18
        out = self.residual_layer(out)
        # 1
        out = self.output(out)
        # 残差
        out = residual+out
        # print(out.shape)
        out = self.last_part(out)
        return out


class Conv_ReLU_Block(nn.Module):

    def __init__(self):
        super(Conv_ReLU_Block, self).__init__()
        self.sequential = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.sequential(x)      

trian_VDSR_ESPCN.py

import argparse, os
import torch
import random
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from vdsr import VDSR_ESPCN
from dataset import DatasetFromVoc

def init():
    # Training settings
    parser = argparse.ArgumentParser(description="the VDSR of Pytorch")

    # batch_size 每次投入模型的图像数据数
    parser.add_argument("--batch_size", type=int, default=32, help="Training batch size")
    # 训练轮次
    parser.add_argument("--epochs", type=int, default=1000, help="Number of epochs to train for")
    # 学习率
    parser.add_argument("--lr", type=float, default=0.0001, help="Learning Rate. Default=0.1")
    # 动态学习率调整系数
    parser.add_argument("--step", type=int, default=10, help="Sets the learning rate to the initial LR decayed by momentum every n epochs, Default: n=10")
    # 使用cuda
    parser.add_argument("--cuda", action="store_true",default=True,help="Use cuda?")
    # 已训练权重
    parser.add_argument("--resume", default="", type=str, help="Path to checkpoint (default: none)")
    # 开始轮次
    parser.add_argument("--start-epoch", default=1, type=int, help="Manual epoch number (useful on restarts)")
    # 梯度裁剪系数
    parser.add_argument("--clip", type=float, default=0.4, help="Clipping Gradients. Default=0.4")
    # 单线程
    parser.add_argument("--num_workers", type=int, default=1, help="Number of threads for data loader to use, Default: 1")
    # 优化器动量
    parser.add_argument("--momentum", default=0.9, type=float, help="Momentum, Default: 0.9")
    # 正则化系数
    parser.add_argument("--weight-decay", "--wd", default=1e-4, type=float, help="Weight decay, Default: 1e-4")
    #  预训练a
    parser.add_argument('--pretrained', default='', type=str, help='path to pretrained model (default: none)')
    # 默认GPU为0
    parser.add_argument("--gpus", default="0", type=str, help="gpu ids (default: 0)")
    return parser

def main():

    parser = init()
    # 获得所有参数
    opt = parser.parse_args()
    print(opt)
    # cuda设置gpu参数
    cuda = opt.cuda

    # gpu配置
    if cuda:
        print("=> use gpu id: '{}'".format(opt.gpus))
        os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
        if not torch.cuda.is_available():
                raise Exception("No GPU found or Wrong gpu id, please run without --cuda")
    # 随机种子参数
    opt.seed = random.randint(1, 10000)
    print("Random Seed: ", opt.seed)
    torch.manual_seed(opt.seed)

    if cuda:
        # 设置固定生成随机数的种子，使得每次运行该.py
        # 文件时生成的随机数相同
        torch.cuda.manual_seed(opt.seed)

    # 设置加速，优化运行效率
    cudnn.benchmark = True

    print("===> Loading datasets")
    train_set = DatasetFromVoc(scale=2)

    # training_data_loader = DataLoader(dataset=train_set,num_workers=opt.num_workers, batch_size=opt.batch_size, shuffle=True)
    training_data_loader = DataLoader(dataset=train_set, batch_size=opt.batch_size,
                                      shuffle=True)
    print("===> Building model")
    model = VDSR_ESPCN(scale=2,input_channel=3)


    criterion = nn.MSELoss()

    print("===> Setting GPU")
    if cuda:
        model = model.cuda()
        criterion = criterion.cuda()

    # optionally resume from a checkpoint
    if opt.resume:
        if os.path.isfile(opt.resume):
            print("=> loading checkpoint '{}'".format(opt.resume))
            checkpoint = torch.load(opt.resume)
            opt.start_epoch = checkpoint["epoch"] + 1
            model.load_state_dict(checkpoint["model"].state_dict())
        else:
            print("=> no checkpoint found at '{}'".format(opt.resume))

    # optionally copy weights from a checkpoint
    if opt.pretrained:

        if os.path.isfile(opt.pretrained):
            print("=> loading model '{}'".format(opt.pretrained))
            weights = torch.load(opt.pretrained)
            model.load_state_dict(weights['model'].state_dict())
        else:
            print("=> no model found at '{}'".format(opt.pretrained))

    print("===> Setting Optimizer")

    optimizer = optim.SGD(model.parameters(), lr=opt.lr, momentum=opt.momentum, weight_decay=opt.weight_decay)

    print("===> Training")

    for epoch in range(opt.start_epoch, opt.epochs + 1):
        train(opt,training_data_loader, optimizer, model, criterion, epoch)
        save_checkpoint(model, epoch)

def adjust_learning_rate(opt, epoch):
    """Sets the learning rate to the initial LR decayed by 10 every 10 epochs"""
    lr = opt.lr * (0.1 ** (epoch //opt.step))
    return lr

def train(opt,training_data_loader, optimizer, model, criterion, epoch):
    # lr = adjust_learning_rate(opt, epoch-1)
    lr = opt.lr
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

    print("Epoch = {}, lr = {}".format(epoch, optimizer.param_groups[0]["lr"]))

    model.train()

    for iteration, batch in enumerate(training_data_loader):
        input, target = torch.Tensor(batch[0]), torch.Tensor(batch[1])

        if opt.cuda:
            input = input.cuda()
            target = target.cuda()

        # print("meodel:",model(input).shape)
        # print(target.shape)
        loss = criterion(model(input), target)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(),opt.clip)
        optimizer.step()

        if iteration%100 == 0:
            print("===> Epoch[{}]({}/{}): Loss: {:.10f}".format(epoch, iteration, len(training_data_loader), loss.item()))

def save_checkpoint(model, epoch):
    model_out_path = "checkpoint_vdsrespcn/" + "model_epoch_{}.pth".format(epoch)
    state = {"epoch": epoch ,"model": model}
    if not os.path.exists("checkpoint_vdsrespcn/"):
        os.makedirs("checkpoint_vdsrespcn/")

    torch.save(state, model_out_path)

    print("Checkpoint saved to {}".format(model_out_path))

if __name__ == "__main__":
    main()      

训练比较久…

• 以上我自己的改写，具体测试还在进行汇总，关于单通道的图片PSNR达到27.65
• 小伙伴们有问题欢迎加微信一起解决哦