ESPCN 论文复现、代码注解、超级详细

环境

文章解释的代码顺序即为代码阅读顺序；
环境：
	ubuntu16.04
	pytorch 1.2
	torchvision 0.40
	cuda 9.2
	python 3.6
	github： [https://github.com/leftthomas/ESPCN](https://github.com/leftthomas/ESPCN)
	上面github应该是官方代码，但是用的pytorch0.4以前的版本应该是，cuda是8.0，python是2.7
	而本文是在官网上的基础上改动了几个函数并翻译，供和我一样的萌新参考。如发现有错误请提醒一下大家共同学习

data_utils.py

import argparse
import os
from os import listdir
from os.path import join

from PIL import Image
from torch.utils.data.dataset import Dataset
"""
    torch.utils.data.Dataset是代表自定义数据集方法的抽象类，你可以自己定义你的数据类继承这个抽象类，非常简单，只需要定义__len__和__getitem__这两个方法就可以。

    通过继承torch.utils.data.Dataset的这个抽象类，我们可以定义好我们需要的数据类。当我们通过迭代的方式来取得每一个数据，但是这样很难实现取batch，shuffle或者多线程读取数据，
    所以pytorch还提供了一个简单的方法来做这件事情，通过torch.utils.data.DataLoader类来定义一个新的迭代器，
    用来将自定义的数据读取接口的输出或者PyTorch已有的数据读取接口的输入按照batch size封装成Tensor，后续只需要再包装成Variable即可作为模型的输入。
    总之，通过torch.utils.data.Dataset和torch.utils.data.DataLoader这两个类，使数据的读取变得非常简单，快捷。

    __len__(self) 定义当被len()函数调用时的行为（返回容器中元素的个数）

    __getitem__(self)定义获取容器中指定元素的行为，相当于self[key]，即允许类对象可以有索引操作

"""
from torchvision.transforms import Compose, CenterCrop, Scale
from tqdm import tqdm


def is_image_file(filename):         #图片文件夹，以及打开格式 
    return any(filename.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.JPG', '.JPEG', '.PNG'])
    """
    endswith():方法用于判断字符串是否以指定后缀结尾，如果以指定后缀结尾返回True，否则返回False。可选参数"start"与"end"为检索字符串的开始与结束位置
    extension: 类扩展
    any():python doc中得说明，意思就是当传入空可迭代对象时返回False，当可迭代对象中有任意一个不为False，则返回True
    即有后面任意一种扩展名都返回  True
    """


def is_video_file(filename):         #视屏文件夹，以及打开格式
    return any(filename.endswith(extension) for extension in ['.mp4', '.avi', '.mpg', '.mkv', '.wmv', '.flv'])


def calculate_valid_crop_size(crop_size, upscale_factor):                   #计算有效切割尺寸
    return crop_size - (crop_size % upscale_factor)


def input_transform(crop_size, upscale_factor):                              #输入图片切割
    return Compose([                                                        #Conpose 用来管理transfrom（数据扩大操作）中的各种类
        CenterCrop(crop_size),                                              #将图片进行中心切割得到 给定size大小图片并用scale缩放
        Scale(crop_size // upscale_factor, interpolation=Image.BICUBIC)     #scale 是缩放   // 代表整除    interpolation是插入类型
    ])


def target_transform(crop_size):                                        #标签图片切割
    return Compose([
        CenterCrop(crop_size)
    ])


class DatasetFromFolder(Dataset):                                         #从文件夹获取数据
    def __init__(self, dataset_dir, upscale_factor, input_transform=None, target_transform=None):
        super(DatasetFromFolder, self).__init__()
        self.image_dir = dataset_dir + '/SRF_' + str(upscale_factor) + '/data'             #路径是 D:\论文代码\ESPCN网络框架\ESPCN-master\data\train\SRF_3\data
        self.target_dir = dataset_dir + '/SRF_' + str(upscale_factor) + '/target'          #路径是D:\论文代码\ESPCN网络框架\ESPCN-master\data\train\SRF_3\target
        #join  连接两个或多个路径名组件，根据需要插入'/'。如果任何组件是绝对路径，则之前的所有路径组件将被丢弃。最后部分为空将生成以分隔符结束的路径
        #listdir == os
        #join   == os.path
        self.image_filenames = [join(self.image_dir, x) for x in listdir(self.image_dir) if is_image_file(x)]   #遍历image_dir文件夹下的文件 ， 符合is_image_file(x)定义格式文件，则把该文件和路径连接
        self.target_filenames = [join(self.target_dir, x) for x in listdir(self.target_dir) if is_image_file(x)]   #target_dir遍历文件夹下的文件 ， 符合is_image_file(x)定义格式文件，则把该文件和路径连接
        self.input_transform = input_transform      #两个空张量，
        self.target_transform = target_transform

    def __getitem__(self, index):
        image, _, _ = Image.open(self.image_filenames[index]).convert('YCbCr').split()  #image.open（）打开图片
        #self.image_filenames([index])是从（ dataset_dir + '/SRF_' + str(upscale_factor) + '/data' ）
        #中索引一张图片用图像处理库 PIL 实现转换其格式为'YCbCr'，并赋值给image

        #PIL的九种不同模式：1(黑白图，二值图)，L（灰度图），P（八色彩图），RGB，RGBA，CMYK（印刷四分色图），YCbCr（模式“YCbCr”为24位彩色图像，Y亮度Cb蓝色色度Cr红色色度）,I，F 

        target, _, _ = Image.open(self.target_filenames[index]).convert('YCbCr').split()  #转换数据图像的格式

        if self.input_transform:
            image = self.input_transform(image)          #将转换后的image 图片切割至给定大小
        if self.target_transform:
            target = self.target_transform(target)       #将转换后的target 图片切割至给定大小

        return image, target

    def __len__(self):
        return len(self.image_filenames)


def generate_dataset(data_type, upscale_factor):
    images_name = [x for x in listdir('data/VOC2012/' + data_type) if is_image_file(x)]  
    #listdir('data/VOC2012/' + data_type)里有扩展名为 ['.mp4', '.avi', '.mpg', '.mkv', '.wmv', '.flv'] 的文件则返回文件并赋值到 image_name 中

    crop_size = calculate_valid_crop_size(256, upscale_factor)  #返回  crop_size = crop_size - (crop_size % upscale_factor)  256 - （256 % 3）= 255   %：求余数
    lr_transform = input_transform(crop_size, upscale_factor)   #把图片中心剪切：crop_size = 255 大小 ； 缩放 crop_size // upscale_factor = 255 // 3 = 0 
    hr_transform = target_transform(crop_size)                  #target 和 input 一样操作

    root = 'data/' + data_type              #data_type是 'data/' 下待创建的： train ; val ; test ;
    if not os.path.exists(root):            #os.makedirs :判断路径存不存在，不i存在创建它
        os.makedirs(root)
    path = root + '/SRF_' + str(upscale_factor)   #合并路径
    if not os.path.exists(path):
        os.makedirs(path)
    image_path = path + '/data'
    if not os.path.exists(image_path):
        os.makedirs(image_path)
    target_path = path + '/target'
    if not os.path.exists(target_path):
        os.makedirs(target_path)
    """
    这一串一直检擦路径存在于否，不存在就创建
    直到创建：D:\论文代码\ESPCN网络框架\ESPCN-master\data        \data_type\SRF_2\target(标签)
           ：D:\论文代码\ESPCN网络框架\ESPCN-master\data        \data_type\SRF_2\data(训练集)
           : data_type是 'data/' 下待创建的： train ; val ; test ;
    """

    for image_name in tqdm(images_name, desc='generate ' + data_type + ' dataset with upscale factor = '+ str(upscale_factor) + ' from VOC2012'):
        #from VOC2012 generate data_type(train,val,test)  dataset with upscale factor = upscale_factor;
        """
        tqdm:是一个快速，可扩展的Python进度条，可以在Python长循环中添加一个进度提示信息，用户只需要封装任意的迭代器tqdm(iterator)即可完成进度条
        #images_name = [x for x in listdir('data/VOC2012/' + data_type) if is_image_file(x)]
        #listdir('data/VOC2012/' + data_type)里有扩展名为 ['.mp4', '.avi', '.mpg', '.mkv', '.wmv', '.flv'] 的文件名则返回文件并赋值到 image_name 中
        """
        image = Image.open('data/VOC2012/' + data_type + '/' + image_name) #将image_name 存的文件名的文件赋值给image
        target = image.copy()
        image = lr_transform(image)             #将VOC2012里面的图片中心切割缩放后放到image 中
        target = hr_transform(target)           #将VOC2012里面的图片中心切割后放到target 中

        image.save(image_path + '/' + image_name)      #把处理后的image  放到  image_path + '/' + image_name  中
        #image_path = \data\data_type\SRF_2\data(训练集)
        target.save(target_path + '/' + image_name)    #把处理后的target  放到 target_path + '/' + image_name  中
        #target_path = \data\data_type\SRF_2\data(训练集)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Generate Super Resolution Dataset')  #用argparse包中ArgumentParser类生成一个参数解析器 parser ；description是描述参数解析器的作用：生成超分辨率数据集
    parser.add_argument('--upscale_factor', default=3, type=int, help='super resolution upscale factor')  #增加一个参数 upscale_factor ,后面是该参数的介绍
    opt = parser.parse_args()                #用对象parse_args()获取参数解析器的参数
    UPSCALE_FACTOR = opt.upscale_factor      #将解析器的参数赋值给 UPSCALE_FACTOR

    generate_dataset(data_type='train', upscale_factor=UPSCALE_FACTOR)    #调用generate_dataset()函数对 VOC2012里面的数据经过转换、剪切、分别制作训练集和验证集
    generate_dataset(data_type='val', upscale_factor=UPSCALE_FACTOR)
    

model.py

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


class Net(nn.Module):
    def __init__(self, upscale_factor):
        super(Net, self).__init__()

        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(32, 1 * (upscale_factor ** 2), (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

    def forward(self, x):
        x = F.tanh(self.conv1(x))
        x = F.tanh(self.conv2(x))
        x = F.sigmoid(self.pixel_shuffle(self.conv3(x)))
        return x


if __name__ == "__main__":
    model = Net(upscale_factor=3)
    print(model)

psnrmeter.py

from math import log10

import torch
from torchnet.meter import meter
"""
meter:提供了一种标准化的方法来测量一系列不同的测量值，这使得测量模型的各种属性变得很容易
三个方法：
    add() 这给 meter 增加了一个观测值
    value() 它将所有的观测结果都考虑在内，返回 meter 的值
    reset() 这样就删除了之前添加的所有观测结果，重新设置了 meter 
"""


class PSNRMeter(meter.Meter):
    def __init__(self):                         #每次调用这个类都刷新一次
        super(PSNRMeter, self).__init__()
        self.reset()

    def reset(self):                            #这样就删除了之前添加的所有观测结果，重新设置了 meter
        self.n = 0
        self.sesum = 0.0

    def add(self, output, target):              #这给 meter 增加了一个观测值
        if not torch.is_tensor(output) and not torch.is_tensor(target):        #确保out target 是张量
            output = torch.from_numpy(output)
            target = torch.from_numpy(target)
        output = output.cpu()                   #cpu计算 变量 output 、 target
        target = target.cpu()
        self.n += output.numel()                #n = n + output.nume1()   n应该是标签数目
        self.sesum += torch.sum((output - target) ** 2)

    def value(self):                            #它将所有的观测结果都考虑在内，返回 meter 的值
        mse = self.sesum / max(1, self.n)
        psnr = 10 * log10(1 / mse)
        return psnr                             #psnr是评价标准： 峰值信噪比

train.py

import argparse

import torch
import torch.nn as nn
import torch.optim as optim
import torchnet as tnt
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch.optim.lr_scheduler import MultiStepLR
from torch.utils.data import DataLoader
from torchnet.engine import Engine
"""
torchnet 是用于 torch 的代码复用和模块化编程的框架
四种方法：
    Dataset : 各种不同的方式处理数据
    Engine: 各种机器学习算法
    Meter: 性能度量评估
    Log: 计算对数
Engine:
    它将训练过程和测试过程进行包装，抽象成一个类，提供train和test方法和一个hooks.(这部分文档是问题的)文档中的描述应该是torch.tensor中的hook，原理一致，
    只不过tensor中的hook是在变量forward或者backward的时候执行（两种hook）
    hooks包括on_start, on_sample, on_forward, on_update,  on_end_epoch,  on_end，可以自己制定函数，在开始，load数据，forward，更新还有epoch结束以及训练结束时执行。
    一般是用开查看和保存模型训练过程的一些结果
"""
from torchnet.logger import VisdomPlotLogger
#用于记录一些评估结果和可视化
from tqdm import tqdm

from data_utils import DatasetFromFolder
from model import Net
from psnrmeter import PSNRMeter


def processor(sample):                   #定义处理器
    """
    processor()函数：
    输入 ：sample是输入的数据集包含训练集数据（data）和标签（target）
    输出 ：预测值 y 和损失值 loss
    """
    data, target, training = sample
    data = Variable(data)               #from torch.autograd import Variable   就是将变量 grad_requirs = True
    target = Variable(target)
    if torch.cuda.is_available():       #检测  GPU  是佛可用
        data = data.cuda()              #将变量放入cuda中
        target = target.cuda()          

    output = model(data)                #调用model  这个data应该是放大因子upscale_factor   output 是输出预测值
    """
        #import torch.nn as nn
        #import torch.nn.functional as F


        #class Net(nn.Module):  
               
          ## torch.nn是专门为神经网络设计的模块化接口。nn构建于autograd之上，可以用来定义和运行神经网络。nn.Module是nn中十分重要的类,包含网络各层的定义及forward方法。
          ##定义自已的网络：
                需要继承nn.Module类，并实现forward方法。
                一般把网络中具有可学习参数的层放在构造函数__init__()中，
                不具有可学习参数的层(如ReLU)可放在构造函数中，也可不放在构造函数中(而在forward中使用nn.functional来代替)
    
                只要在nn.Module的子类中定义了forward函数，backward函数就会被自动实现(利用Autograd)。
                在forward函数中可以使用任何Variable支持的函数，毕竟在整个pytorch构建的图中，是Variable在流动。还可以使用
                if,for,print,log等python语法.
    
                注：Pytorch基于nn.Module构建的模型中，只支持mini-batch的Variable输入方式，
                比如，只有一张输入图片，也需要变成 N x C x H x W 的形式：
    
                input_image = torch.FloatTensor(1, 28, 28)
                input_image = Variable(input_image)
                input_image = input_image.unsqueeze(0)   # 1 x 1 x 28 x 28

            def __init__(self, upscale_factor):                                      ##upscale_factor  :上采样因素 就是放大倍数
                super(Net, self).__init__()

                self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))                ##self.conv2d = nn.Conv2d(in_channels=3,out_channels=64,kernel_size=4,stride=2,padding=1)
                self.conv2 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
                self.conv3 = nn.Conv2d(32, 1 * (upscale_factor ** 2), (3, 3), (1, 1), (1, 1))     ##此处的upscale_factor 就是为后续的上采样实现  
                self.pixel_shuffle = nn.PixelShuffle(upscale_factor)                 ##PixelShuffle  亚像素实现上采样： input:(m , c*upscale_factor**2 , h , w)
            #                                                                                                         oueput:(m , c , h*upscale_facter , w*upscale_facter)

            def forward(self, x):                                                     ##  nn.Module实现的层（layer）是一个特殊的类，都是由class Layer(nn.Module)定义，会自动提取可学习的参数
            #                                                                         ##  nn.functional中的函数更像是纯函数,由def functional(input）定义  一般定义不可学习参数的层，如激活、池化
                x = F.tanh(self.conv1(x))              ##定义激活函数 import torch.nn.functional as F
                x = F.tanh(self.conv2(x))
                x = F.sigmoid(self.pixel_shuffle(self.conv3(x)))
                return x


        if __name__ == "__main__":
            model = Net(upscale_factor=3)
            print(model)
    """

    loss = criterion(output, target)      #调用损失函数  一般用法：  criterion = LossCriterion()
    #                                      output是预测值     target是标签真实值             loss = criterion(y, Y)

    return loss, output


def on_sample(state):
    state['sample'].append(state['train'])         #给 train 这个键赋值


def reset_meters():
    meter_psnr.reset()         #清空：评价标准 和 损失数据 
    meter_loss.reset()


def on_forward(state):
    meter_psnr.add(state['output'].data, state['sample'][1])   #计算评价标准：tensor.data 切断反向传播，因为评价标准不要学习参数（可用tensor.detach()更好） add（）：把output 和 sample[1] 元素逐个相加
    meter_loss.add(state['loss'].item())                       #.team()返回一个可遍历的数组 ；计算的损失逐个相加


def on_start_epoch(state):
    reset_meters()               #清空：评价标准 和 损失数据 

    scheduler.step()      #在scheduler的step_size表示scheduler.step()每调用step_size次，对应的学习率就会按照策略调整一次。
    #所以如果scheduler.step()是放在mini-batch里面，那么step_size指的是经过这么多次迭代，学习率改变一次

    state['iterator'] = tqdm(state['iterator']) #iterator迭代器，返回迭代进度条提示信息
    #tqdm: 是一个快速，可扩展的Python进度条，可以在Python长循环中添加一个进度提示信息，用户只需要封装任意的迭代器tqdm(iterator)即可完成进度条


def on_end_epoch(state):
    print('[Epoch %d] Train Loss: %.4f (PSNR: %.2f db)' % (state['epoch'], meter_loss.value()[0], meter_psnr.value()))  #打印训练的迭代次数对应的 损失 峰值信噪比

    train_loss_logger.log(state['epoch'], meter_loss.value()[0])      #计算训练集损失和峰值信噪比的对数，以迭代次数为底
    train_psnr_logger.log(state['epoch'], meter_psnr.value())

    reset_meters()                     #清空：评价标准 和 损失数据

    engine.test(processor, val_loader)    #用engine 封装test的方法， processor（）函数是向前传播计算预测值和损失； 
    #后面定义   ：val_loader = DataLoader(dataset=val_set, num_workers=4, batch_size=64, shuffle=False)

    val_loss_logger.log(state['epoch'], meter_loss.value()[0])         #计算验证集损失和峰值信噪比的对数，以迭代次数为底
    val_psnr_logger.log(state['epoch'], meter_psnr.value())

    print('[Epoch %d] Val Loss: %.4f (PSNR: %.2f db)' % (state['epoch'], meter_loss.value()[0], meter_psnr.value()))    #打印验证的迭代次数对应的 损失 峰值信噪比

    torch.save(model.state_dict(), 'epochs/epoch_%d_%d.pt' % (UPSCALE_FACTOR, state['epoch']))  #state_dict()：返回一个包含模块的整个状态的字典（返回model字典里是：前向卷积和激活函数）
    #save（）：将对象（返回的model字典）保存到磁盘文件中。


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Train Super Resolution') ##用argparse包中ArgumentParser类生成一个参数解析器 parser ；description是描述参数解析器的作用：训练超分
    parser.add_argument('--upscale_factor', default=3, type=int, help='super resolution upscale factor')  #增加一个  upscale_factor参数 ：超分上采样因数
    parser.add_argument('--num_epochs', default=100, type=int, help='super resolution epochs number')   #增加一个  num_epochs参数：超分迭代次数
    opt = parser.parse_args()        ##用对象parse_args()获取参数解析器的参数到 opt里

    UPSCALE_FACTOR = opt.upscale_factor
    NUM_EPOCHS = opt.num_epochs

    train_set = DatasetFromFolder('data/train', upscale_factor=UPSCALE_FACTOR, input_transform=transforms.ToTensor(),target_transform=transforms.ToTensor())
    #调用 data_utils.py 里的类 DatasetFromFolder：对数据进行处理、转换、剪切 得到 image(训练用的data) 和 target（标签# ） 后面赋值给dataset 进行封装batch_size 进行训练
    
    val_set = DatasetFromFolder('data/val', upscale_factor=UPSCALE_FACTOR, input_transform=transforms.ToTensor(),target_transform=transforms.ToTensor())

    train_loader = DataLoader(dataset=train_set, num_workers=4, batch_size=64, shuffle=True)
    """
    先简单的介绍一下DataLoader
    它是PyTorch中数据读取的一个重要接口，该接口定义在dataloader.py中，只要是用PyTorch来训练模型基本都会用到该接口（除非用户重写…），
    该接口的目的：将自定义的Dataset根据batch size大小、是否shuffle等封装成一个Batch Size大小的Tensor，用于后面的训练。

    # torch.utils.data.DataLoader() 将自定义的Dataset根据batch size大小、是否shuffle等封装成一个Batch Size大小的Tensor，用于后面的训练。
       torch.utils.data.DataLoader（
                            dataset，#数据加载
                            batch_size = 1，#批处理大小设置
                            shuffle = False，#是否进项洗牌操作
                            sampler = None，#指定数据加载中使用的索引/键的序列
                            batch_sampler = None，#和sampler类似
                            num_workers = 0，#是否进行多进程加载数据设置
                            collat​​e_fn = None，#是否合并样本列表以形成一小批Tensor
                            pin_memory = False，#如果True，数据加载器会在返回之前将Tensors复制到CUDA固定内存
                            drop_last = False，#True如果数据集大小不能被批处理大小整除，则设置为删除最后一个不完整的批处理。
                            timeout = 0，#如果为正，则为从工作人员收集批处理的超时值
                            worker_init_fn = None ）

    """
    val_loader = DataLoader(dataset=val_set, num_workers=4, batch_size=64, shuffle=False)   #对训练集和验证集设置  batch_size 批量处理

    model = Net(upscale_factor=UPSCALE_FACTOR)  #前向传播
    criterion = nn.MSELoss()                  #计算损失
    if torch.cuda.is_available():           #把model 里的计算添加到cuda 中计算
        model = model.cuda()
        criterion = criterion.cuda()

    print('# parameters:', sum(param.numel() for param in model.parameters()))

    optimizer = optim.Adam(model.parameters(), lr=1e-2)                         #优化方法为adma；  model.parameters :待优化的参数 ； lr：学习率
    scheduler = MultiStepLR(optimizer, milestones=[30, 80], gamma=0.1)          #学习率变化：milestones表示学习率在 30 80 时候各变化一次 变化 gamma = 0.1 倍

    engine = Engine()          #封装test train val 的训练过程
    meter_loss = tnt.meter.AverageValueMeter()  #计算平均回归损失 
    meter_psnr = PSNRMeter()  #获取峰值信噪比

    train_loss_logger = VisdomPlotLogger('line', opts={'title': 'Train Loss'})    #可视化
    train_psnr_logger = VisdomPlotLogger('line', opts={'title': 'Train PSNR'})
    val_loss_logger = VisdomPlotLogger('line', opts={'title': 'Val Loss'})
    val_psnr_logger = VisdomPlotLogger('line', opts={'title': 'Val PSNR'})

    engine.hooks['on_sample'] = on_sample                                         #engine封装的hook，即训练的过程文档
    engine.hooks['on_forward'] = on_forward
    engine.hooks['on_start_epoch'] = on_start_epoch
    engine.hooks['on_end_epoch'] = on_end_epoch

    engine.train(processor, train_loader, maxepoch=NUM_EPOCHS, optimizer=optimizer) #engine封装的训练

结束！图片测试和视频测试都和 train.py 大同小异 供大家自己解读；
萌新第一次发博客求点赞 ：）

本文地址：https://blog.csdn.net/qq_38477064/article/details/109641687