本文对代码进行了详细的注释
论文地址：https://arxiv.org/abs/1803.02735
代码地址：https://github.com/alterzero/DBPN-Pytorch
base_networks
import torch
import math
class DenseBlock(torch.nn.Module):
def __init__(self, input_size, output_size, bias=True, activation='relu', norm='batch'):
super(DenseBlock, self).__init__()
self.fc = torch.nn.Linear(input_size, output_size, bias=bias) #nn.Linear()的初始化：
self.norm = norm #normalization的意思，即是否要进行标准化
if self.norm =='batch':
self.bn = torch.nn.BatchNorm1d(output_size)
#nn.BatchNorm1d 是 PyTorch 中的一个用于一维数据（例如序列或时间序列）的批标准化（Batch Normalization）层。
elif self.norm == 'instance':
self.bn = torch.nn.InstanceNorm1d(output_size)
#torch.nn.InstanceNorm1d是 PyTorch 中的一个用于一维实例归一化：对一个批次中每个样本，依次按照通道计算对应的均值及均方差。
#选择不同的激活函数：
self.activation = activation
if self.activation == 'relu':
self.act = torch.nn.ReLU(True)
elif self.activation == 'prelu':
self.act = torch.nn.PReLU()
elif self.activation == 'lrelu':
self.act = torch.nn.LeakyReLU(0.2, True)
elif self.activation == 'tanh':
self.act = torch.nn.Tanh()
elif self.activation == 'sigmoid':
self.act = torch.nn.Sigmoid()
def forward(self, x):
if self.norm is not None:
#如果有进行批量归一化的提示（满足第10行或13行），因此需要先变线性，然后再经过33行的self.bn函数，根据提示（满足第10行还是13行），选择进行批量归一化或者实例归一化。
out = self.bn(self.fc(x))
else:
out = self.fc(x)#否则（也就是self.norm is None：没有说明要归一化），则直接只经过线性处理
if self.activation is not None:
return self.act(out) #如果需要选择激活函数，则给out选择激活函数，将结果返回出来
else:
return out #否则不选择激活函数，直接返回out
class ConvBlock(torch.nn.Module):
def __init__(self, input_size, output_size, kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm=None):
super(ConvBlock, self).__init__()
self.conv = torch.nn.Conv2d(input_size, output_size, kernel_size, stride, padding, bias=bias)
self.norm = norm
if self.norm =='batch':
self.bn = torch.nn.BatchNorm2d(output_size)
elif self.norm == 'instance':
self.bn = torch.nn.InstanceNorm2d(output_size)
self.activation = activation
if self.activation == 'relu':
self.act = torch.nn.ReLU(True)
elif self.activation == 'prelu':
self.act = torch.nn.PReLU()
elif self.activation == 'lrelu':
self.act = torch.nn.LeakyReLU(0.2, True)
elif self.activation == 'tanh':
self.act = torch.nn.Tanh()
elif self.activation == 'sigmoid':
self.act = torch.nn.Sigmoid()
def forward(self, x):
if self.norm is not None:
out = self.bn(self.conv(x))#因为是ConvBlock，所以就在这里先用卷积处理
else:
out = self.conv(x)
if self.activation is not None:
return self.act(out)
else:
return out
class DeconvBlock(torch.nn.Module):
def __init__(self, input_size, output_size, kernel_size=4, stride=2, padding=1, bias=True, activation='prelu', norm=None):
super(DeconvBlock, self).__init__()
self.deconv = torch.nn.ConvTranspose2d(input_size, output_size, kernel_size, stride, padding, bias=bias)
self.norm = norm
if self.norm == 'batch':
self.bn = torch.nn.BatchNorm2d(output_size)
elif self.norm == 'instance':
self.bn = torch.nn.InstanceNorm2d(output_size)
self.activation = activation
if self.activation == 'relu':
self.act = torch.nn.ReLU(True)
elif self.activation == 'prelu':
self.act = torch.nn.PReLU()
elif self.activation == 'lrelu':
self.act = torch.nn.LeakyReLU(0.2, True)
elif self.activation == 'tanh':
self.act = torch.nn.Tanh()
elif self.activation == 'sigmoid':
self.act = torch.nn.Sigmoid()
def forward(self, x):
if self.norm is not None:
out = self.bn(self.deconv(x))#因为是DeconvBlock，所以就在这里先用反卷积处理
else:
out = self.deconv(x)
if self.activation is not None:
return self.act(out)
else:
return out
class ResnetBlock(torch.nn.Module):
def __init__(self, num_filter, kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm='batch'):
super(ResnetBlock, self).__init__()
self.conv1 = torch.nn.Conv2d(num_filter, num_filter, kernel_size, stride, padding, bias=bias)
self.conv2 = torch.nn.Conv2d(num_filter, num_filter, kernel_size, stride, padding, bias=bias)
self.norm = norm
if self.norm == 'batch':
self.bn = torch.nn.BatchNorm2d(num_filter)
elif norm == 'instance':
self.bn = torch.nn.InstanceNorm2d(num_filter)
#不同的激活函数：
self.activation = activation
if self.activation == 'relu':
self.act = torch.nn.ReLU(True)
elif self.activation == 'prelu':
self.act = torch.nn.PReLU()
elif self.activation == 'lrelu':
self.act = torch.nn.LeakyReLU(0.2, True)
elif self.activation == 'tanh':
self.act = torch.nn.Tanh()
elif self.activation == 'sigmoid':
self.act = torch.nn.Sigmoid()
def forward(self, x):
residual = x #残差也就是要反投影的误差为x
if self.norm is not None:
out = self.bn(self.conv1(x))#在文章流程图中我们能看出，需要用公式（4）将残差特征图的尺寸提高到和H^0t 相同，这样才能用公式（5）将两个特征图相加合成输出的特征图Ht
else:
out = self.conv1(x)#所以这里面也是必须要进行卷积
if self.activation is not None:
out = self.act(out)
if self.norm is not None:
out = self.bn(self.conv2(out))#这个out也要进行卷积，实际上这个out是低分辨率的特征图L^t-1
else:
out = self.conv2(out)
out = torch.add(out, residual)#把res和out（H^0t）算在一起，这里的add就是投影单元的那个加法器
return out
class UpBlock(torch.nn.Module):#是在写明上投影单元中的一个最开始的那个基础的stage（如161-163行所示）是怎么具体的实现上投影的作用的（即文章中的Fig3上面的图片的具体流程用代码写出来）这是非密集的上投影单元
#这个UpBlock，放在dbpn.py中看的出来，作用是提供一个（非密集的）上投影单元的操作
def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, bias=True, activation='prelu', norm=None):
super(UpBlock, self).__init__()
self.up_conv1 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.up_conv2 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.up_conv3 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
def forward(self, x):
h0 = self.up_conv1(x)#x相当于论文中Fig3中的Lt-1，h0相当于H0t，
l0 = self.up_conv2(h0)#l0相当于L0t
h1 = self.up_conv3(l0 - x)#？？为什么l0-x：见上投影单元的具体结构，实际上是求出了残差e tl这个东西，然后这一行实际上是将l0-x送进up_conv3，而up_conv3对应的是DeconvBlock模块，正好和上投影单元的第二个蓝色方块对应上了，即公式（3）
return h1 + h0
class UpBlockPix(torch.nn.Module):#没见到这个函数在那里用到了？？？？？！
def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, scale=4, bias=True, activation='prelu', norm=None):
super(UpBlockPix, self).__init__()
self.up_conv1 = Upsampler(scale,num_filter)
self.up_conv2 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.up_conv3 = Upsampler(scale,num_filter)
def forward(self, x):
h0 = self.up_conv1(x)
l0 = self.up_conv2(h0)
h1 = self.up_conv3(l0 - x)
return h1 + h0
class D_UpBlock(torch.nn.Module):#是在写明上投影单元是怎么具体的实现上投影的作用的（即文章中的Fig4上面的图片的具体流程用代码写出来）这是密集的上投影单元
def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, num_stages=1, bias=True, activation='prelu', norm=None):
super(D_UpBlock, self).__init__()
self.conv = ConvBlock(num_filter*num_stages, num_filter, 1, 1, 0, activation, norm=None)#密集的上投影单元这里多了一个conv(1x1),它体现在了Fig4的流程图之中
self.up_conv1 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.up_conv2 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.up_conv3 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
def forward(self, x):
x = self.conv(x)#密集的上投影单元这里多了一个conv(1x1),它体现在了Fig4的流程图之中
h0 = self.up_conv1(x)#和前面的166行同理
l0 = self.up_conv2(h0)
h1 = self.up_conv3(l0 - x)
return h1 + h0
class D_UpBlockPix(torch.nn.Module):#也没见过使用它
def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, num_stages=1, scale=4, bias=True, activation='prelu', norm=None):
super(D_UpBlockPix, self).__init__()
self.conv = ConvBlock(num_filter*num_stages, num_filter, 1, 1, 0, activation, norm=None)
self.up_conv1 = Upsampler(scale,num_filter)
self.up_conv2 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.up_conv3 = Upsampler(scale,num_filter)
def forward(self, x):
x = self.conv(x)
h0 = self.up_conv1(x)
l0 = self.up_conv2(h0)
h1 = self.up_conv3(l0 - x)
return h1 + h0
class DownBlock(torch.nn.Module):
def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, bias=True, activation='prelu', norm=None):
super(DownBlock, self).__init__()
self.down_conv1 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.down_conv2 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.down_conv3 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
def forward(self, x):#和前面166行同理
l0 = self.down_conv1(x)
h0 = self.down_conv2(l0)
l1 = self.down_conv3(h0 - x)
return l1 + l0
class DownBlockPix(torch.nn.Module):
def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, scale=4,bias=True, activation='prelu', norm=None):
super(DownBlockPix, self).__init__()
self.down_conv1 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.down_conv2 = Upsampler(scale,num_filter)
self.down_conv3 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
def forward(self, x):
l0 = self.down_conv1(x)
h0 = self.down_conv2(l0)
l1 = self.down_conv3(h0 - x)
return l1 + l0
class D_DownBlock(torch.nn.Module):
def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, num_stages=1, bias=True, activation='prelu', norm=None):
super(D_DownBlock, self).__init__()
self.conv = ConvBlock(num_filter*num_stages, num_filter, 1, 1, 0, activation, norm=None)#也是先conv(1x1)
self.down_conv1 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.down_conv2 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.down_conv3 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
def forward(self, x):
x = self.conv(x)
l0 = self.down_conv1(x)
h0 = self.down_conv2(l0)
l1 = self.down_conv3(h0 - x)
return l1 + l0
class D_DownBlockPix(torch.nn.Module):
def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, num_stages=1, scale=4, bias=True, activation='prelu', norm=None):
super(D_DownBlockPix, self).__init__()
self.conv = ConvBlock(num_filter*num_stages, num_filter, 1, 1, 0, activation, norm=None)
self.down_conv1 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
self.down_conv2 = Upsampler(scale,num_filter)
self.down_conv3 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
def forward(self, x):
x = self.conv(x)
l0 = self.down_conv1(x)
h0 = self.down_conv2(l0)
l1 = self.down_conv3(h0 - x)
return l1 + l0
class PSBlock(torch.nn.Module):#没有发现在哪里用到了，是为了下面347行调用写的
def __init__(self, input_size, output_size, scale_factor, kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm='batch'):
super(PSBlock, self).__init__()
self.conv = torch.nn.Conv2d(input_size, output_size * scale_factor**2, kernel_size, stride, padding, bias=bias)
self.ps = torch.nn.PixelShuffle(scale_factor)
self.norm = norm
if self.norm == 'batch':
self.bn = torch.nn.BatchNorm2d(output_size)
elif norm == 'instance':
self.bn = torch.nn.InstanceNorm2d(output_size)
self.activation = activation
if self.activation == 'relu':
self.act = torch.nn.ReLU(True)
elif self.activation == 'prelu':
self.act = torch.nn.PReLU()
elif self.activation == 'lrelu':
self.act = torch.nn.LeakyReLU(0.2, True)
elif self.activation == 'tanh':
self.act = torch.nn.Tanh()
elif self.activation == 'sigmoid':
self.act = torch.nn.Sigmoid()
def forward(self, x):
if self.norm is not None:
out = self.bn(self.ps(self.conv(x)))
else:
out = self.ps(self.conv(x))
#上面几行和下面几行都是在和之前进行一样的操作，上面是标准化normalize，下面是激活函数activation
if self.activation is not None:
out = self.act(out)
return out
class Upsampler(torch.nn.Module):#不清楚，再观察一下？？？?
def __init__(self, scale, n_feat, bn=False, act='prelu', bias=True):
super(Upsampler, self).__init__()
modules = []
for _ in range(int(math.log(scale, 2))):
modules.append(ConvBlock(n_feat, 4 * n_feat, 3, 1, 1, bias, activation=None, norm=None))
modules.append(torch.nn.PixelShuffle(2))
if bn: modules.append(torch.nn.BatchNorm2d(n_feat))
#modules.append(torch.nn.PReLU())
self.up = torch.nn.Sequential(*modules)
self.activation = act
if self.activation == 'relu':
self.act = torch.nn.ReLU(True)
elif self.activation == 'prelu':
self.act = torch.nn.PReLU()
elif self.activation == 'lrelu':
self.act = torch.nn.LeakyReLU(0.2, True)
elif self.activation == 'tanh':
self.act = torch.nn.Tanh()
elif self.activation == 'sigmoid':
self.act = torch.nn.Sigmoid()
def forward(self, x):
out = self.up(x)
if self.activation is not None:
out = self.act(out)
return out
class Upsample2xBlock(torch.nn.Module):#不清楚，再观察一下？？？?
def __init__(self, input_size, output_size, bias=True, upsample='deconv', activation='relu', norm='batch'):
super(Upsample2xBlock, self).__init__()
scale_factor = 2
# 1. Deconvolution (Transposed convolution)
if upsample == 'deconv':
self.upsample = DeconvBlock(input_size, output_size,
kernel_size=4, stride=2, padding=1,
bias=bias, activation=activation, norm=norm)
# 2. Sub-pixel convolution (Pixel shuffler)
elif upsample == 'ps':
self.upsample = PSBlock(input_size, output_size, scale_factor=scale_factor,
bias=bias, activation=activation, norm=norm)
# 3. Resize and Convolution
elif upsample == 'rnc':
self.upsample = torch.nn.Sequential(
torch.nn.Upsample(scale_factor=scale_factor, mode='nearest'),
ConvBlock(input_size, output_size,
kernel_size=3, stride=1, padding=1,
bias=bias, activation=activation, norm=norm)
)
def forward(self, x):
out = self.upsample(x)
return out
eval.py
from __future__ import print_function
import argparse
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
from dbpn import Net as DBPN
from dbpn_v1 import Net as DBPNLL
from dbpn_iterative import Net as DBPNITER
from data import get_eval_set
from functools import reduce
#from scipy.misc import imsave
import scipy.io as sio
import time
import cv2
import imageio
# Training settings，下面都是定义的一些参数
parser = argparse.ArgumentParser(description='PyTorch Super Res Example')
parser.add_argument('--upscale_factor', type=int, default=2, help="super resolution upscale factor")
parser.add_argument('--testBatchSize', type=int, default=1, help='testing batch size')#batchsize是一批样本的大小。将整个样本分成若干个batch，每个batch中样本的数量就叫做batchsize。
parser.add_argument('--gpu_mode', type=bool, default=True)
parser.add_argument('--self_ensemble', type=bool, default=False)
parser.add_argument('--chop_forward', type=bool, default=False)
parser.add_argument('--threads', type=int, default=0, help='number of threads for data loader to use')
parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123')
parser.add_argument('--gpus', default=1, type=int, help='number of gpu')
parser.add_argument('--input_dir', type=str, default='Input')
parser.add_argument('--output', default='Results/', help='Location to save checkpoint models')
parser.add_argument('--test_dataset', type=str, default='Set5_LR_x2')
#这里面31和33行是相互关联的，也就是说实际上实验的测试数据在D:\experiments\DBPN-Pytorch-master\Input\DIV2K_valid_HR之中，但是运行本文件时说不可以，说GPU不够！我把测试集改成了最基础的Set5_LR_x2是可以的。
parser.add_argument('--model_type', type=str, default='DBPN')#现在是选择了DBPN这个模型类型，具体是选择哪个模型需要看36行（因为有不同的scale factor）
parser.add_argument('--residual', type=bool, default=False)
parser.add_argument('--model', default='models/DBPN_x2.pth', help='sr pretrained base model')#具体选择的模型
opt = parser.parse_args()
gpus_list=range(opt.gpus)
print(opt)
cuda = opt.gpu_mode
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
test_set = get_eval_set(os.path.join(opt.input_dir,opt.test_dataset), opt.upscale_factor)
testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.testBatchSize, shuffle=False)
print('===> Building model')
if opt.model_type == 'DBPNLL':
model = DBPNLL(num_channels=3, base_filter=64, feat = 256, num_stages=10, scale_factor=opt.upscale_factor) ###D-DBPN
elif opt.model_type == 'DBPN-RES-MR64-3':
model = DBPNITER(num_channels=3, base_filter=64, feat = 256, num_stages=3, scale_factor=opt.upscale_factor) ###D-DBPN
else:#这种情况说明在用DBPN这种模型
model = DBPN(num_channels=3, base_filter=64, feat = 256, num_stages=7, scale_factor=opt.upscale_factor) ###D-DBPN
if cuda:
model = torch.nn.DataParallel(model, device_ids=gpus_list)
model.load_state_dict(torch.load(opt.model, map_location=lambda storage, loc: storage))
print('Pre-trained SR model is loaded.')
if cuda:
model = model.cuda(gpus_list[0])
def eval():
model.eval()
for batch in testing_data_loader:
with torch.no_grad():
input, bicubic, name = Variable(batch[0]), Variable(batch[1]), batch[2]
if cuda:
input = input.cuda(gpus_list[0]) #转到GPU上面运行
bicubic = bicubic.cuda(gpus_list[0])
t0 = time.time()
if opt.chop_forward:
with torch.no_grad():
prediction = chop_forward(input, model, opt.upscale_factor)
else:
if opt.self_ensemble:
with torch.no_grad():
prediction = x8_forward(input, model)
else:
with torch.no_grad():
prediction = model(input)
if opt.residual:
prediction = prediction + bicubic
t1 = time.time()
print("===> Processing: %s || Timer: %.4f sec." % (name[0], (t1 - t0)))
save_img(prediction.cpu().data, name[0])
def save_img(img, img_name):
save_img = img.squeeze().clamp(0, 1).numpy().transpose(1,2,0)
# save img
save_dir=os.path.join(opt.output,opt.test_dataset) #https://wenku.csdn.net/answer/b8f850e9e6b14671b9e4eaed967a69f6
if not os.path.exists(save_dir):
os.makedirs(save_dir)
save_fn = save_dir +'/'+ img_name
cv2.imwrite(save_fn, cv2.cvtColor(save_img*255, cv2.COLOR_BGR2RGB), [cv2.IMWRITE_PNG_COMPRESSION, 0])
def x8_forward(img, model, precision='single'):
def _transform(v, op):
if precision != 'single': v = v.float()
v2np = v.data.cpu().numpy()
if op == 'vflip':
tfnp = v2np[:, :, :, ::-1].copy()
elif op == 'hflip':
tfnp = v2np[:, :, ::-1, :].copy()
elif op == 'transpose':
tfnp = v2np.transpose((0, 1, 3, 2)).copy()
ret = torch.Tensor(tfnp).cuda()
if precision == 'half':
ret = ret.half()
elif precision == 'double':
ret = ret.double()
with torch.no_grad():
ret = Variable(ret)
return ret
inputlist = [img]
for tf in 'vflip', 'hflip', 'transpose':
inputlist.extend([_transform(t, tf) for t in inputlist])
outputlist = [model(aug) for aug in inputlist]
for i in range(len(outputlist)):
if i > 3:
outputlist[i] = _transform(outputlist[i], 'transpose')
if i % 4 > 1:
outputlist[i] = _transform(outputlist[i], 'hflip')
if (i % 4) % 2 == 1:
outputlist[i] = _transform(outputlist[i], 'vflip')
output = reduce((lambda x, y: x + y), outputlist) / len(outputlist)
return output
def chop_forward(x, model, scale, shave=8, min_size=80000, nGPUs=opt.gpus):
b, c, h, w = x.size()
h_half, w_half = h // 2, w // 2
h_size, w_size = h_half + shave, w_half + shave
inputlist = [
x[:, :, 0:h_size, 0:w_size],
x[:, :, 0:h_size, (w - w_size):w],
x[:, :, (h - h_size):h, 0:w_size],
x[:, :, (h - h_size):h, (w - w_size):w]]
if w_size * h_size < min_size:
outputlist = []
for i in range(0, 4, nGPUs):
with torch.no_grad():
input_batch = torch.cat(inputlist[i:(i + nGPUs)], dim=0)
if opt.self_ensemble:
with torch.no_grad():
output_batch = x8_forward(input_batch, model)
else:
with torch.no_grad():
output_batch = model(input_batch)
outputlist.extend(output_batch.chunk(nGPUs, dim=0))
else:
outputlist = [
chop_forward(patch, model, scale, shave, min_size, nGPUs) \
for patch in inputlist]
h, w = scale * h, scale * w
h_half, w_half = scale * h_half, scale * w_half
h_size, w_size = scale * h_size, scale * w_size
shave *= scale
with torch.no_grad():
output = Variable(x.data.new(b, c, h, w))
output[:, :, 0:h_half, 0:w_half] \
= outputlist[0][:, :, 0:h_half, 0:w_half]
output[:, :, 0:h_half, w_half:w] \
= outputlist[1][:, :, 0:h_half, (w_size - w + w_half):w_size]
output[:, :, h_half:h, 0:w_half] \
= outputlist[2][:, :, (h_size - h + h_half):h_size, 0:w_half]
output[:, :, h_half:h, w_half:w] \
= outputlist[3][:, :, (h_size - h + h_half):h_size, (w_size - w + w_half):w_size]
return output
##Eval Start!!!!
eval()
main.py
from __future__ import print_function
import argparse
from math import log10
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from dbpn import Net as DBPN #Back-projection stages中up scale有7个
from dbpn_v1 import Net as DBPNLL #dbpn_v1中的DBPNLL中 Back-projection stages中up scale比DBPN（7个）多了三个（10个）
from dbpns import Net as DBPNS #dbpns中的DBPNS中 Back-projection stages中up scale只有2个。DBPN（7个）(本实验没用到！！！)
from dbpn_iterative import Net as DBPNITER #和迭代相关，但具体情况没搞清楚呢
from data import get_training_set
import pdb
import socket
import time
# Training settings
parser = argparse.ArgumentParser(description='PyTorch Super Res Example')
parser.add_argument('--upscale_factor', type=int, default=2, help="super resolution upscale factor")
parser.add_argument('--batchSize', type=int, default=1, help='training batch size')#batchsize=1
parser.add_argument('--nEpochs', type=int, default=20, help='number of epochs to train for')#epochs=20，
# 已知一个epoch是将整个训练样本训练模型一次，那为什么需要多个epoch（训练多次）？因为单次训练是不够的，需要反复多次训练才能拟合收敛。每个epoch都在更新神经网络中的权重参数。所以epoch太多或者太少都不好
parser.add_argument('--snapshots', type=int, default=5, help='Snapshots') #snapshots：快照版本，即每 5 epochs保存一次，则保存次数=epochs/snapshots
parser.add_argument('--start_iter', type=int, default=1, help='Starting Epoch')
parser.add_argument('--lr', type=float, default=1e-4, help='Learning Rate. Default=0.01')#学习率
parser.add_argument('--gpu_mode', type=bool, default=True)
parser.add_argument('--threads', type=int, default=0, help='number of threads for data loader to use')
parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123')
parser.add_argument('--gpus', default=1, type=int, help='number of gpu')
parser.add_argument('--data_dir', type=str, default='D:\experiments\DBPN-Pytorch-master\Dataset')#数据来源根目录
parser.add_argument('--data_augmentation', type=bool, default=True)
parser.add_argument('--hr_train_dataset', type=str, default='DIV2K_train_HR')#具体选择的是哪个训练集呢？是DIV2K_train_HR
parser.add_argument('--train_dataset', type=str, default='DIV2K_train_HR')
parser.add_argument('--model_type', type=str, default='DBPN')#选择模型
parser.add_argument('--residual', type=bool, default=True)
parser.add_argument('--patch_size', type=int, default=20, help='Size of cropped HR image')#图像块被分成多大的小块
parser.add_argument('--pretrained_sr', default='MIX2K_LR_aug_x4dl10DBPNITERtpami_epoch_399.pth', help='sr pretrained base model')
parser.add_argument('--pretrained', type=bool, default=False)
parser.add_argument('--save_folder', default='weights/', help='Location to save checkpoint models')#folder是文件夹的意思，
# 那为什么这里的路径不用写前面具体的路径呢？我猜是因为在anaconda prompt中，已经确定了具体的路径在这个master文件夹里面执行了，所以只需要写到这里就可以了
parser.add_argument('--prefix', default='tpami_residual_filter8', help='Location to save checkpoint models')
opt = parser.parse_args()
gpus_list = range(opt.gpus)
hostname = str(socket.gethostname())
cudnn.benchmark = True
print(opt)
def train(epoch):
epoch_loss = 0
model.train()
for iteration, batch in enumerate(training_data_loader, 1):
input, target, bicubic = Variable(batch[0]), Variable(batch[1]), Variable(batch[2])
if cuda:
input = input.cuda(gpus_list[0])
target = target.cuda(gpus_list[0])
bicubic = bicubic.cuda(gpus_list[0])
optimizer.zero_grad()
t0 = time.time()
prediction = model(input)
if opt.residual:
prediction = prediction + bicubic
loss = criterion(prediction, target)#这里是在用criterion（）函数求损失
t1 = time.time()
epoch_loss += loss.data #！！！+= 代表原位操作（见总结）
loss.backward()
optimizer.step()
print("===> Epoch[{}]({}/{}): Loss: {:.4f} || Timer: {:.4f} sec.".format(epoch, iteration, len(training_data_loader), loss.data, (t1 - t0)))
print("===> Epoch {} Complete: Avg. Loss: {:.4f}".format(epoch, epoch_loss / len(training_data_loader)))
def test():
avg_psnr = 0
for batch in testing_data_loader:
input, target = Variable(batch[0]), Variable(batch[1])
if cuda:
input = input.cuda(gpus_list[0])
target = target.cuda(gpus_list[0])
prediction = model(input)
mse = criterion(prediction, target)#criterion（）函数是MSE，即本文当中的均方误差
psnr = 10 * log10(1 / mse.data[0])#求PSNR
avg_psnr += psnr
print("===> Avg. PSNR: {:.4f} dB".format(avg_psnr / len(testing_data_loader)))
def print_network(net):
num_params = 0
for param in net.parameters():
num_params += param.numel()
print(net)
print('Total number of parameters: %d' % num_params)
def checkpoint(epoch):
model_out_path = opt.save_folder+opt.train_dataset+hostname+opt.model_type+opt.prefix+"_epoch_{}.pth".format(epoch)
torch.save(model.state_dict(), model_out_path)
print("Checkpoint saved to {}".format(model_out_path))
cuda = opt.gpu_mode
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
print('===> Loading datasets')
train_set = get_training_set(opt.data_dir, opt.hr_train_dataset, opt.upscale_factor, opt.patch_size, opt.data_augmentation)
training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batchSize, shuffle=True)
print('===> Building model ', opt.model_type)
if opt.model_type == 'DBPNLL': #这种情况，就是指DBPNLL
model = DBPNLL(num_channels=3, base_filter=64, feat = 256, num_stages=10, scale_factor=opt.upscale_factor)
elif opt.model_type == 'DBPN-RES-MR64-3': #这种情况，就是指DBPNITER
model = DBPNITER(num_channels=3, base_filter=64, feat = 256, num_stages=3, scale_factor=opt.upscale_factor)
else: #这种否则的情况，就是指DBPN
model = DBPN(num_channels=3, base_filter=64, feat = 256, num_stages=7, scale_factor=opt.upscale_factor)
model = torch.nn.DataParallel(model, device_ids=gpus_list)
criterion = nn.L1Loss()
print('---------- Networks architecture -------------')
print_network(model)
print('----------------------------------------------')
if opt.pretrained:
model_name = os.path.join(opt.save_folder + opt.pretrained_sr)
if os.path.exists(model_name):
#model= torch.load(model_name, map_location=lambda storage, loc: storage)
model.load_state_dict(torch.load(model_name, map_location=lambda storage, loc: storage))
print('Pre-trained SR model is loaded.')
if cuda:
model = model.cuda(gpus_list[0])
criterion = criterion.cuda(gpus_list[0])
optimizer = optim.Adam(model.parameters(), lr=opt.lr, betas=(0.9, 0.999), eps=1e-8)
for epoch in range(opt.start_iter, opt.nEpochs + 1):
train(epoch)
# learning rate is decayed by a factor of 10 every half of total epochs
if (epoch+1) % (opt.nEpochs/2) == 0:
for param_group in optimizer.param_groups:
param_group['lr'] /= 10.0
print('Learning rate decay: lr={}'.format(optimizer.param_groups[0]['lr']))
if (epoch+1) % (opt.snapshots) == 0:
checkpoint(epoch)
dbpn.py
import os
import torch.nn as nn
import torch.optim as optim
from base_networks import *
from torchvision.transforms import *
class Net(nn.Module):
def __init__(self, num_channels, base_filter, feat, num_stages, scale_factor): #__init __: 用于向类中传入外部参数，同时定义样本集
super(Net, self).__init__()
if scale_factor == 2:
kernel = 6
stride = 2
padding = 2
elif scale_factor == 4:
kernel = 8
stride = 4
padding = 2
elif scale_factor == 8:
kernel = 12
stride = 8
padding = 2
#Initial Feature Extraction
#kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm=None
self.feat0 = ConvBlock(num_channels, feat, 3, 1, 1, activation='prelu', norm=None)
self.feat1 = ConvBlock(feat, base_filter, 1, 1, 0, activation='prelu', norm=None)#先按照流程图的特征提取模块，进行两次卷积
#Back-projection stages：stages=7
#七个上投影单元，六个下投影单元，所以stage是7（一对上下采样级为一个stage）
self.up1 = UpBlock(base_filter, kernel, stride, padding)
self.down1 = DownBlock(base_filter, kernel, stride, padding)
self.up2 = UpBlock(base_filter, kernel, stride, padding)
self.down2 = D_DownBlock(base_filter, kernel, stride, padding, 2)
self.up3 = D_UpBlock(base_filter, kernel, stride, padding, 2) #D_UpBlock表示密集投影，上面那几行不用密集投影是因为他们前面没有可以用来作为密集投影输入的那么多个单元
self.down3 = D_DownBlock(base_filter, kernel, stride, padding, 3)
self.up4 = D_UpBlock(base_filter, kernel, stride, padding, 3)
self.down4 = D_DownBlock(base_filter, kernel, stride, padding, 4)
self.up5 = D_UpBlock(base_filter, kernel, stride, padding, 4)
self.down5 = D_DownBlock(base_filter, kernel, stride, padding, 5)
self.up6 = D_UpBlock(base_filter, kernel, stride, padding, 5)
self.down6 = D_DownBlock(base_filter, kernel, stride, padding, 6)
self.up7 = D_UpBlock(base_filter, kernel, stride, padding, 6)
#Reconstruction
self.output_conv = ConvBlock(num_stages*base_filter, num_channels, 3, 1, 1, activation=None, norm=None)
#模型参数初始化方法：自定义初始化函数：
for m in self.modules():
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)#何凯明等提出的一种鲁棒的神经网络参数初始化方法，是为了保证信息在前向传播和反向传播过程中能够有效流动，使不同层的输入信号的方差大致相等。
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('ConvTranspose2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x):
x = self.feat0(x)
x = self.feat1(x)
h1 = self.up1(x) #h是high，l是low
l1 = self.down1(h1) #up是上采样得到的结果
h2 = self.up2(l1)
#上面这三行是刚开始的一个stage，因为是第一个，所以涉及不到串联的问题
# 下面这三行是后面的stage，所以涉及了串联的问题
concat_h = torch.cat((h2, h1),1) #(h2, h1)说明要把h2和h1连在一起，作为下投影unit的输入
l = self.down2(concat_h)
concat_l = torch.cat((l, l1),1)
h = self.up3(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down3(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up4(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down4(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up5(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down5(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up6(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down6(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up7(concat_l)
concat_h = torch.cat((h, concat_h),1)#concat_h
x = self.output_conv(concat_h) #送入到重建模块进行重建，得到最终的结果x
# self.output_conv是重建函数，因此得到的最终结果应该是x
#注意混淆：密集连接（dense connections）是所有之前projection（投影）的输入都进入后面投影级做输出；
# concatenation是指所有的上投影单元得到的（各自包含了不同的特征）HR图像，将他们都串联起来后作为重建模块的输入
return x
dbpns.py
import os
import torch.nn as nn
import torch.optim as optim
from base_networks import *
from torchvision.transforms import *
class Net(nn.Module):
def __init__(self, num_channels, base_filter, feat, num_stages, scale_factor):
super(Net, self).__init__()
#关于上一行super的一些解释：https://blog.csdn.net/qq_31244453/article/details/104657532?ops_request_misc=&request_id=&biz_id=102&utm_term=super(Net,%20self).__init__()&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-1-104657532.142^v96^pc_search_result_base9&spm=1018.2226.3001.4187
if scale_factor == 2:
kernel = 6
stride = 2
padding = 2
elif scale_factor == 4:
kernel = 8
stride = 4
padding = 2
elif scale_factor == 8:
kernel = 12
stride = 8
padding = 2
#Initial Feature Extraction
self.feat0 = ConvBlock(num_channels, feat, 3, 1, 1, activation='prelu', norm=None)
self.feat1 = ConvBlock(feat, base_filter, 1, 1, 0, activation='prelu', norm=None)
#Back-projection stages
#应该是DBPN，因为没有密集投影（应该是文章说的那个lighter的那个网络SS network，但是本实验没用它）
self.up1 = UpBlock(base_filter, kernel, stride, padding)
self.down1 = DownBlock(base_filter, kernel, stride, padding)
self.up2 = UpBlock(base_filter, kernel, stride, padding)
#Reconstruction
self.output_conv = ConvBlock(num_stages*base_filter, num_channels, 3, 1, 1, activation=None, norm=None)
for m in self.modules():#自定义初始化参数
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('ConvTranspose2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x):
x = self.feat0(x)
x = self.feat1(x)
h1 = self.up1(x)
h2 = self.up2(self.down1(h1))
x = self.output_conv(torch.cat((h2, h1),1))
return x
dbpn-v1.py
import os
import torch.nn as nn
import torch.optim as optim
from base_networks import *
from torchvision.transforms import *
class Net(nn.Module):
def __init__(self, num_channels, base_filter, feat, num_stages, scale_factor):
super(Net, self).__init__()
if scale_factor == 2:
kernel = 6
stride = 2
padding = 2
elif scale_factor == 4:
kernel = 8
stride = 4
padding = 2
elif scale_factor == 8:
kernel = 12
stride = 8
padding = 2
#Initial Feature Extraction
self.feat0 = ConvBlock(num_channels, feat, 3, 1, 1, activation='prelu', norm=None)
self.feat1 = ConvBlock(feat, base_filter, 1, 1, 0, activation='prelu', norm=None)
#Back-projection stages：stages=10
# 10个上投影单元，9个下投影单元，所以stage是10
self.up1 = UpBlock(base_filter, kernel, stride, padding)
self.down1 = DownBlock(base_filter, kernel, stride, padding)
self.up2 = UpBlock(base_filter, kernel, stride, padding)
self.down2 = D_DownBlock(base_filter, kernel, stride, padding, 2)#也是密集投影了（有dense）
self.up3 = D_UpBlock(base_filter, kernel, stride, padding, 2)
self.down3 = D_DownBlock(base_filter, kernel, stride, padding, 3)
self.up4 = D_UpBlock(base_filter, kernel, stride, padding, 3)
self.down4 = D_DownBlock(base_filter, kernel, stride, padding, 4)
self.up5 = D_UpBlock(base_filter, kernel, stride, padding, 4)
self.down5 = D_DownBlock(base_filter, kernel, stride, padding, 5)
self.up6 = D_UpBlock(base_filter, kernel, stride, padding, 5)
self.down6 = D_DownBlock(base_filter, kernel, stride, padding, 6)
self.up7 = D_UpBlock(base_filter, kernel, stride, padding, 6)
self.down7 = D_DownBlock(base_filter, kernel, stride, padding, 7)
self.up8 = D_UpBlock(base_filter, kernel, stride, padding, 7)
self.down8 = D_DownBlock(base_filter, kernel, stride, padding, 8)
self.up9 = D_UpBlock(base_filter, kernel, stride, padding, 8)
self.down9 = D_DownBlock(base_filter, kernel, stride, padding, 9)
self.up10 = D_UpBlock(base_filter, kernel, stride, padding, 9)
#Reconstruction
self.output_conv = ConvBlock(num_stages*base_filter, num_channels, 3, 1, 1, activation=None, norm=None)
# 模型参数初始化方法：自定义初始化函数：
for m in self.modules():
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('ConvTranspose2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x):
x = self.feat0(x)
x = self.feat1(x)
h1 = self.up1(x)
l1 = self.down1(h1)
h2 = self.up2(l1)
concat_h = torch.cat((h2, h1),1) #这一看就知道这是一个D-DBPN的过程
l = self.down2(concat_h)
concat_l = torch.cat((l, l1),1)
h = self.up3(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down3(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up4(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down4(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up5(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down5(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up6(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down6(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up7(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down7(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up8(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down8(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up9(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down9(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up10(concat_l) #从这里也能看出，stages=10,因为是将各个HR串联
concat_h = torch.cat((h, concat_h),1)#这里等式右侧的h是最后一次upprojection产生的HR10，concat_h是指之前九次的HR的串联好的结果。在这行中，将两者串联在一起，就可以得到所有HR图像的串联结果（即反向投影模块的结果）
x = self.output_conv(concat_h)#重建，得到最终结果
return x
dbpn-iterative.py
import os
import torch.nn as nn
import torch.optim as optim
from base_networks import *
from torchvision.transforms import *
class Net(nn.Module):
def __init__(self, num_channels, base_filter, feat, num_stages, scale_factor):
super(Net, self).__init__()
if scale_factor == 2:
kernel = 6
stride = 2
padding = 2
elif scale_factor == 4:
kernel = 8
stride = 4
padding = 2
elif scale_factor == 8:
kernel = 12
stride = 8
padding = 2
self.num_stages = num_stages
#Initial Feature Extraction
self.feat0 = ConvBlock(num_channels, feat, 3, 1, 1, activation='prelu', norm=None)
self.feat1 = ConvBlock(feat, base_filter, 1, 1, 0, activation='prelu', norm=None)
#Back-projection stages
self.up1 = UpBlock(base_filter, kernel, stride, padding)
self.down1 = DownBlock(base_filter, kernel, stride, padding)
self.up2 = UpBlock(base_filter, kernel, stride, padding)
self.down2 = D_DownBlock(base_filter, kernel, stride, padding, 2)
self.up3 = D_UpBlock(base_filter, kernel, stride, padding, 2)
self.down3 = D_DownBlock(base_filter, kernel, stride, padding, 3)
self.up4 = D_UpBlock(base_filter, kernel, stride, padding, 3)
self.down4 = D_DownBlock(base_filter, kernel, stride, padding, 4)
self.up5 = D_UpBlock(base_filter, kernel, stride, padding, 4)
self.down5 = D_DownBlock(base_filter, kernel, stride, padding, 5)
self.up6 = D_UpBlock(base_filter, kernel, stride, padding, 5)
self.down6 = D_DownBlock(base_filter, kernel, stride, padding, 6)
self.up7 = D_UpBlock(base_filter, kernel, stride, padding, 6)
#Reconstruction
self.output_conv = ConvBlock(num_stages*base_filter, num_channels, 3, 1, 1, activation=None, norm=None)
for m in self.modules():
classname = m.__class__.__name__
if classname.find('Conv2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('ConvTranspose2d') != -1:
torch.nn.init.kaiming_normal_(m.weight)
if m.bias is not None:
m.bias.data.zero_()
def forward(self, x):
x = self.feat0(x)
l = self.feat1(x)
results = []
for i in range(self.num_stages):
#它和dbpn的区别是这个是迭代循环！也就是说循环次数是和num_stages相关联的。
# 而且根据for i in range这个结构，是将从1-（range-1）依次赋给i，和下面的append相配合，一起构成整个循环后串联的工作。
h1 = self.up1(l)
l1 = self.down1(h1)
h2 = self.up2(l1)
concat_h = torch.cat((h2, h1),1)
l = self.down2(concat_h)
concat_l = torch.cat((l, l1),1)
h = self.up3(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down3(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up4(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down4(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up5(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down5(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up6(concat_l)
concat_h = torch.cat((h, concat_h),1)
l = self.down6(concat_h)
concat_l = torch.cat((l, concat_l),1)
h = self.up7(concat_l)
results.append(h)#append的作用就是将最后一个求出的h加到数组中，因为这里是一个循环结构，循环次数正好和stages数和上投影单元数（T）相对应，所以详见append的作用
results = torch.cat(results,1)#当上面的循环结构已经将得到的HR全部串联在一起后，就得到了反投影的结果。
x = self.output_conv(results)
return x
扫一扫
723
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
