版权声明： https://blog.csdn.net/gwplovekimi/article/details/85774605
目录
python train_sub.py -opt options/train/train_noise_blur_sr.json
tensorboard --logdir tb_logger/ --port 6008
处理数据的代码可以参考本人的GitHub（https://github.com/gwpscut/degradation-model-for-image-restoration）
setting
{
"name": "noiseestimation_blurestimation_SR" // please remove "debug_" during training
, "tb_logger_dir": "sr_noise_blur"
, "use_tb_logger": true
, "model":"sr_noise_blur"
, "scale": 4
, "crop_scale": 4
, "gpu_ids": [3,4]
// , "init_type": "kaiming"
//
// , "finetune_type": "basic" //sft | basic
, "datasets": {
"train": {
"name": "DIV2K"
, "mode": "LRMRMATHR"
, "dataroot_HR": "/home/guanwp/BasicSR_datasets/DIV2K800_sub"
, "dataroot_MR": "/home/guanwp/BasicSR_datasets/DIV2K800_sub_blur_bicLRx4"//the target for the noise estimation
, "dataroot_LR": "/home/guanwp/BasicSR_datasets/DIV2K800_sub_blur_bicLRx4_noiseALL"
, "dataroot_MAT": "/home/guanwp/BasicSR_datasets/DIV2K800_sub_estimation"//the target for the blur estimation
, "subset_file": null
, "use_shuffle": true
, "n_workers": 8
, "batch_size": 24 // 16
, "HR_size": 128 // 128 | 192 | 96
, "noise_gt": true//residual for the noise
, "use_flip": true
, "use_rot": true
}
, "val": {
"name": "val_set5_x4_c03s08_mod4",
"mode": "LRHR",
"dataroot_HR": "/home/guanwp/BasicSR_datasets/val_set5/Set5",
"dataroot_LR": "/home/guanwp/BasicSR_datasets/val_set5/Set5_blur_bicLRx4_noiseALL"
}
}
, "path": {
"root": "/home/guanwp/Blind_Restoration-master/sr_noise_blur"
// , "pretrain_model_G": null
// , "pretrain_model_sub_noise": null
// , "pretrain_model_sub_blur": null
}
, "network_G": {
"which_model_G": "sr_resnet" // sr_resnet | modulate_sr_resnet
// , "norm_type": "sft"
, "norm_type": null
, "mode": "CNA"
, "nf": 64
, "nb": 16
, "in_nc": 9
, "out_nc": 3
// , "gc": 32
, "group": 1
// , "gate_conv_bias": true
}
////
, "network_sub": {
"which_model_sub": "noise_subnet" // sr_resnet |noise_subnet
// , "norm_type": "adaptive_conv_res"
, "norm_type": "batch"
// , "norm_type": null
, "mode": "CNA"
, "nf": 64
// , "nb": 16
, "in_nc": 3
, "out_nc": 3
, "group": 1
// , "down_scale": 2
}
, "network_sub2": {
"which_model_sub": "blur_subnet" // sr_resnet | blur_subnet
// , "norm_type": "adaptive_conv_res"
, "norm_type": "batch"
// , "norm_type": null
, "mode": "CNA"
, "nf": 64
// , "nb": 16
, "in_nc": 6
, "out_nc": 3
, "group": 1
// , "down_scale": 2
}
, "train": {
// "lr_G": 1e-3
"lr_G": 1e-4
, "lr_scheme": "MultiStepLR"
// , "lr_steps": [200000, 400000, 600000, 800000]
, "lr_steps": [500000]
// , "lr_steps": [600000]
// , "lr_steps": [1000000]
, "lr_gamma": 0.1
// , "lr_gamma": 0.5
, "pixel_criterion_basic": "l2"
, "pixel_criterion_noise": "l2"
, "pixel_criterion_reg_noise": "tv"
, "pixel_criterion_blur": "l2"
, "pixel_criterion_reg_blur": "tv"
, "pixel_weight_basic": 1.0
, "pixel_weight_noise": 1.0
, "pixel_weight_blur": 1.0
, "val_freq": 1e3
, "manual_seed": 0
, "niter": 1e6
// , "niter": 6e5
}
, "logger": {
"print_freq": 200
, "save_checkpoint_freq": 1e3
}
}
数据处理中的.mat文件
LRMRMATHR_dataset.py
import os.path
import random
import numpy as np
import cv2
import torch
import torch.utils.data as data
import data.util as util
from scipy.io import loadmat
class LRMRMATHRDataset(data.Dataset):
'''
Read LR, MR and HR image pair.
The pair is ensured by 'sorted' function, so please check the name convention.
'''
def __init__(self, opt):
super(LRMRMATHRDataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.paths_MR = None
self.paths_HR = None
self.paths_MAT = None
self.LR_env = None # environment for lmdb
self.MR_env = None
self.HR_env = None
self.MAT_env = None
self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_HR'])
self.MR_env, self.paths_MR = util.get_image_paths(opt['data_type'], opt['dataroot_MR'])
self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
self.MAT_env, self.paths_MAT = util.get_image_paths(opt['data_type'], opt['dataroot_MAT'])
assert self.paths_HR, 'Error: HR path is empty.'
if self.paths_LR and self.paths_MR:
assert len(self.paths_LR) == len(self.paths_MR), \
'MR and LR datasets have different number of images - {}, {}.'.format(\
len(self.paths_LR), len(self.paths_MR))
self.random_scale_list = [1]
def __getitem__(self, index):
HR_path, LR_path, MR_path, MAT_path = None, None, None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# get HR image
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
# # modcrop in the validation / test phase
# if self.opt['phase'] != 'train':
# img_HR = util.modcrop(img_HR, scale)
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
MR_path = self.paths_MR[index]
img_MR = util.read_img(self.MR_env, MR_path)
# get mat file
MAT_path = self.paths_MAT[index]
img_MAT = loadmat(MAT_path)['im_residual']
# kernel_gt = loadmat(MAT_path)['kernel_gt']
# img_MAT = np.zeros_like(img_LR)
if self.opt['noise_gt']:
img_MR = img_LR - img_MR
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, C = img_LR.shape
LR_size = HR_size // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
img_MR = img_MR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
img_MAT = img_MAT[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
# for ind, value in enumerate(kernel_gt):
# img_MAT[:, :, ind] = np.tile(value, (LR_size, LR_size))
# augmentation - flip, rotate
img_MR, img_MAT, img_LR, img_HR = util.augment([img_MR, img_MAT, img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_MR = img_MR[:, :, [2, 1, 0]]
img_MAT = img_MAT[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
img_MR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_MR, (2, 0, 1)))).float()
img_MAT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_MAT, (2, 0, 1)))).float()
return {'HR': img_HR, 'LR': img_LR, 'MR': img_MR, 'MAT': img_MAT, 'HR_path': HR_path, 'MR_path': MR_path,
'LR_path': LR_path, 'MAT_path': MAT_path}
def __len__(self):
return len(self.paths_HR)
LRMRHR_dataset.py
import os.path
import random
import numpy as np
import cv2
import torch
import torch.utils.data as data
import data.util as util
class LRMRHRDataset(data.Dataset):
'''
Read LR, MR and HR image pair.
The pair is ensured by 'sorted' function, so please check the name convention.
'''
def __init__(self, opt):
super(LRMRHRDataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.paths_MR = None
self.paths_HR = None
self.LR_env = None # environment for lmdb
self.MR_env = None
self.HR_env = None
self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_HR'])
self.MR_env, self.paths_MR = util.get_image_paths(opt['data_type'], opt['dataroot_MR'])
self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
assert self.paths_HR, 'Error: HR path is empty.'
if self.paths_LR and self.paths_MR:
assert len(self.paths_LR) == len(self.paths_MR), \
'MR and LR datasets have different number of images - {}, {}.'.format(\
len(self.paths_LR), len(self.paths_MR))
self.random_scale_list = [1]
def __getitem__(self, index):
HR_path, LR_path, MR_path = None, None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# get HR image
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
# modcrop in the validation / test phase
# if self.opt['phase'] != 'train':
# img_HR = util.modcrop(img_HR, scale)
# change color space if necessary
if self.opt['color']:
img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'], [img_HR])[0]
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
MR_path = self.paths_MR[index]
img_MR = util.read_img(self.MR_env, MR_path)
if self.opt['noise_gt']:
img_MR = img_LR - img_MR
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, C = img_LR.shape
LR_size = HR_size // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
img_MR = img_MR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
# augmentation - flip, rotate
img_MR, img_LR, img_HR = util.augment([img_MR, img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# channel conversion
if self.opt['color']:
# img_HR, img_LR, img_MR = util.channel_convert(C, self.opt['color'], [img_HR, img_LR, img_MR])
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
img_MR = util.channel_convert(C, self.opt['color'], [img_MR])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_MR = img_MR[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
img_MR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_MR, (2, 0, 1)))).float()
return {'HR': img_HR, 'LR': img_LR, 'MR': img_MR, 'HR_path': HR_path, 'MR_path': MR_path, 'LR_path': LR_path}
def __len__(self):
return len(self.paths_HR)
model
关键部分就是model结构的设计。需要到各网络的输出contact到一起
import os
from collections import OrderedDict
import torch
import torch.nn as nn
from torch.optim import lr_scheduler
import models.networks as networks
from .base_model import BaseModel
from .modules.loss import TVLoss
class SRModel(BaseModel):
def __init__(self, opt):
super(SRModel, self).__init__(opt)
train_opt = opt['train']
finetune_type = opt['finetune_type']
# define network and load pretrained models
self.netG = networks.define_G(opt).to(self.device)
self.subnet_noise = networks.define_sub(opt).to(self.device)
self.subnet_blur = networks.define_sub2(opt).to(self.device)
self.load()
if self.is_train:
self.netG.train()
if finetune_type in ['basic', 'sft_basic', 'sft', 'sub_sft']:
self.subnet_noise.eval()
self.subnet_blur.eval()
else:
self.subnet_noise.train()
self.subnet_blur.train()
# loss on noise
loss_type_noise = train_opt['pixel_criterion_noise']
if loss_type_noise == 'l1':
self.cri_pix_noise = nn.L1Loss().to(self.device)
elif loss_type_noise == 'l2':
self.cri_pix_noise = nn.MSELoss().to(self.device)
else:
raise NotImplementedError('Noise loss type [{:s}] is not recognized.'.format(loss_type_noise))
self.l_pix_noise_w = train_opt['pixel_weight_noise']
loss_reg_noise = train_opt['pixel_criterion_reg_noise']
if loss_reg_noise == 'tv':
self.cri_pix_reg_noise = TVLoss(0.00001).to(self.device)
# loss on blur
loss_type_blur = train_opt['pixel_criterion_blur']
if loss_type_blur == 'l1':
self.cri_pix_blur = nn.L1Loss().to(self.device)
elif loss_type_blur == 'l2':
self.cri_pix_blur = nn.MSELoss().to(self.device)
else:
raise NotImplementedError('Blur loss type [{:s}] is not recognized.'.format(loss_type_blur))
self.l_pix_blur_w = train_opt['pixel_weight_blur']
loss_reg_blur = train_opt['pixel_criterion_reg_blur']
if loss_reg_blur == 'tv':
self.cri_pix_reg_blur = TVLoss(0.00001).to(self.device)
loss_type_basic = train_opt['pixel_criterion_basic']
if loss_type_basic == 'l1':
self.cri_pix_basic = nn.L1Loss().to(self.device)
elif loss_type_basic == 'l2':
self.cri_pix_basic = nn.MSELoss().to(self.device)
else:
raise NotImplementedError('Loss type [{:s}] is not recognized.'.format(loss_type_basic))
self.l_pix_basic_w = train_opt['pixel_weight_basic']
# optimizers
wd_G = train_opt['weight_decay_G'] if train_opt['weight_decay_G'] else 0
self.optim_params = self.__define_grad_params(finetune_type)
self.optimizer_G = torch.optim.Adam(
self.optim_params, lr=train_opt['lr_G'], weight_decay=wd_G)
self.optimizers.append(self.optimizer_G)
# schedulers
if train_opt['lr_scheme'] == 'MultiStepLR':
for optimizer in self.optimizers:
self.schedulers.append(lr_scheduler.MultiStepLR(optimizer, \
train_opt['lr_steps'], train_opt['lr_gamma']))
else:
raise NotImplementedError('MultiStepLR learning rate scheme is enough.')
self.log_dict = OrderedDict()
print('---------- Model initialized ------------------')
self.print_network()
print('-----------------------------------------------')
def feed_data(self, data, need_MR=True, need_MAT=True):
self.var_L = data['LR'].to(self.device) # LR
self.real_H = data['HR'].to(self.device) # HR
if need_MR:
self.mid_L = data['MR'].to(self.device) # MR
if need_MAT:
self.real_blur = data['MAT'].to(self.device)
def __define_grad_params(self, finetune_type=None):
optim_params = []
if finetune_type == 'sft':
for k, v in self.netG.named_parameters():
v.requires_grad = False
if k.find('Gate') >= 0:
v.requires_grad = True
optim_params.append(v)
print('we only optimize params: {}'.format(k))
else:
for k, v in self.netG.named_parameters(): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
print('params [{:s}] will optimize.'.format(k))
else:
print('WARNING: params [{:s}] will not optimize.'.format(k))
for k, v in self.subnet_noise.named_parameters(): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
print('params [{:s}] will optimize.'.format(k))
else:
print('WARNING: params [{:s}] will not optimize.'.format(k))
for k, v in self.subnet_blur.named_parameters(): # can optimize for a part of the model
if v.requires_grad:
optim_params.append(v)
print('params [{:s}] will optimize.'.format(k))
else:
print('WARNING: params [{:s}] will not optimize.'.format(k))
return optim_params
def optimize_parameters(self, step):
self.optimizer_G.zero_grad()
self.fake_noise = self.subnet_noise(self.var_L)
l_pix_noise = self.l_pix_noise_w * self.cri_pix_noise(self.fake_noise, self.mid_L)
l_pix_noise = l_pix_noise + self.cri_pix_reg_noise(self.fake_noise)
input_noise = torch.cat((self.var_L, self.fake_noise), 1)
self.fake_blur = self.subnet_blur(input_noise)
l_pix_blur = self.l_pix_blur_w * self.cri_pix_blur(self.fake_blur*16, self.real_blur)
l_pix_blur = l_pix_blur + self.cri_pix_reg_blur(self.fake_blur)
input_noise_blur = torch.cat((input_noise, self.fake_blur), 1)
self.fake_H = self.netG(input_noise_blur)
l_pix_basic = self.l_pix_basic_w * self.cri_pix_basic(self.fake_H, self.real_H)
l_pix = l_pix_noise + l_pix_blur + l_pix_basic
l_pix.backward()
self.optimizer_G.step()
self.log_dict['l_pix'] = l_pix.item()
def test(self):
self.netG.eval()
self.subnet_noise.eval()
self.subnet_blur.eval()
if self.is_train:
for v in self.optim_params:
v.requires_grad = False
else:
for k, v in self.netG.named_parameters():
v.requires_grad = False
for k, v in self.subnet_noise.named_parameters():
v.requires_grad = False
for k, v in self.subnet_blur.named_parameters():
v.requires_grad = False
self.fake_noise = self.subnet_noise(self.var_L)
input_noise = torch.cat((self.var_L, self.fake_noise), 1)
self.fake_blur = self.subnet_blur(input_noise)
input_noise_blur = torch.cat((input_noise, self.fake_blur), 1)
self.fake_H = self.netG(input_noise_blur)
if self.is_train:
for v in self.optim_params:
v.requires_grad = True
else:
for k, v in self.netG.named_parameters():
v.requires_grad = True
for k, v in self.subnet_noise.named_parameters():
v.requires_grad = True
for k, v in self.subnet_blur.named_parameters():
v.requires_grad = True
self.netG.train()
if self.opt['finetune_type'] in ['basic', 'sft_basic', 'sft', 'sub_sft']:
self.subnet_noise.eval()
self.subnet_blur.eval()
else:
self.subnet_noise.train()
self.subnet_blur.eval()
# def test(self):
# self.netG.eval()
# for k, v in self.netG.named_parameters():
# v.requires_grad = False
# self.fake_H = self.netG(self.var_L)
# for k, v in self.netG.named_parameters():
# v.requires_grad = True
# self.netG.train()
def get_current_log(self):
return self.log_dict
def get_current_visuals(self, need_HR=True):
out_dict = OrderedDict()
out_dict['LR'] = self.var_L.detach()[0].float().cpu()
out_dict['MR'] = self.fake_noise.detach()[0].float().cpu()
out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
if need_HR:
out_dict['HR'] = self.real_H.detach()[0].float().cpu()
return out_dict
def print_network(self):
# G
s, n = self.get_network_description(self.netG)
print('Number of parameters in G: {:,d}'.format(n))
if self.is_train:
message = '-------------- Generator --------------\n' + s + '\n'
network_path = os.path.join(self.save_dir, '../', 'network.txt')
with open(network_path, 'w') as f:
f.write(message)
# noise subnet
s, n = self.get_network_description(self.subnet_noise)
print('Number of parameters in noise subnet: {:,d}'.format(n))
message = '\n\n\n-------------- noise subnet --------------\n' + s + '\n'
with open(network_path, 'a') as f:
f.write(message)
# blur subnet
s, n = self.get_network_description(self.subnet_blur)
print('Number of parameters in blur subnet: {:,d}'.format(n))
message = '\n\n\n-------------- blur subnet --------------\n' + s + '\n'
with open(network_path, 'a') as f:
f.write(message)
def load(self):
load_path_G = self.opt['path']['pretrain_model_G']
load_path_sub_noise = self.opt['path']['pretrain_model_sub_noise']
load_path_sub_blur = self.opt['path']['pretrain_model_sub_blur']
if load_path_G is not None:
print('loading model for G [{:s}] ...'.format(load_path_G))
self.load_network(load_path_G, self.netG)
if load_path_sub_noise is not None:
print('loading model for noise subnet [{:s}] ...'.format(load_path_sub_noise))
self.load_network(load_path_sub_noise, self.subnet_noise)
if load_path_sub_blur is not None:
print('loading model for blur subnet [{:s}] ...'.format(load_path_sub_blur))
self.load_network(load_path_sub_blur, self.subnet_blur)
def save(self, iter_label):
self.save_network(self.save_dir, self.netG, 'G', iter_label)
self.save_network(self.save_dir, self.subnet_noise, 'sub_noise', iter_label)
self.save_network(self.save_dir, self.subnet_blur, 'sub_blur', iter_label)
network
至于网络的结构，blur和noise estimation subnetwork都是采用DNCNN的结构，而SR网络采用srresnet
在network中需要定义两个subnetwork
import functools
import torch
import torch.nn as nn
from torch.nn import init
import models.modules.architecture as arch
import models.modules.sft_arch as sft_arch
####################
# initialize
####################
def weights_init_normal(m, std=0.02):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
init.normal_(m.weight.data, 0.0, std)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('Linear') != -1:
init.normal_(m.weight.data, 0.0, std)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('BatchNorm2d') != -1:
init.normal_(m.weight.data, 1.0, std) # BN also uses norm
init.constant_(m.bias.data, 0.0)
def weights_init_kaiming(m, scale=1):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
m.weight.data *= scale
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('Linear') != -1:
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
m.weight.data *= scale
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('BatchNorm2d') != -1 or classname.find('InstanceNorm2d') != -1:
init.constant_(m.weight.data, 1.0)
init.constant_(m.bias.data, 0.0)
# elif classname.find('AdaptiveConvResNorm') != -1:
# init.constant_(m.weight.data, 0.0)
# if m.bias is not None:
# m.bias.data.zero_()
def weights_init_orthogonal(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
init.orthogonal_(m.weight.data, gain=1)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('Linear') != -1:
init.orthogonal_(m.weight.data, gain=1)
if m.bias is not None:
m.bias.data.zero_()
elif classname.find('BatchNorm2d') != -1:
init.constant_(m.weight.data, 1.0)
init.constant_(m.bias.data, 0.0)
def init_weights(net, init_type='kaiming', scale=1, std=0.02):
# scale for 'kaiming', std for 'normal'.
print('initialization method [{:s}]'.format(init_type))
if init_type == 'normal':
weights_init_normal_ = functools.partial(weights_init_normal, std=std)
net.apply(weights_init_normal_)
elif init_type == 'kaiming':
weights_init_kaiming_ = functools.partial(weights_init_kaiming, scale=scale)
net.apply(weights_init_kaiming_)
elif init_type == 'orthogonal':
net.apply(weights_init_orthogonal)
else:
raise NotImplementedError('initialization method [{:s}] not implemented'.format(init_type))
####################
# define network
####################
# Generator
def define_G(opt):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_G']
which_model = opt_net['which_model_G']
if which_model == 'sr_resnet': # SRResNet
netG = arch.SRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], \
nb=opt_net['nb'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'], \
act_type='relu', mode=opt_net['mode'], upsample_mode='pixelshuffle')
elif which_model == 'modulate_sr_resnet':
netG = arch.ModulateSRResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
upscale=opt_net['scale'], norm_type=opt_net['norm_type'], mode=opt_net['mode'],
upsample_mode='pixelshuffle', ada_ksize=opt_net['ada_ksize'],
gate_conv_bias=opt_net['gate_conv_bias'])
elif which_model == 'arcnn':
netG = arch.ARCNN(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'],
norm_type=opt_net['norm_type'], mode=opt_net['mode'], ada_ksize=opt_net['ada_ksize'])
elif which_model == 'srcnn':
netG = arch.SRCNN(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'],
norm_type=opt_net['norm_type'], mode=opt_net['mode'], ada_ksize=opt_net['ada_ksize'])
elif which_model == 'denoise_resnet':
netG = arch.DenoiseResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
upscale=opt_net['scale'], norm_type=opt_net['norm_type'], mode=opt_net['mode'],
upsample_mode='pixelshuffle', ada_ksize=opt_net['ada_ksize'],
down_scale=opt_net['down_scale'], fea_norm=opt_net['fea_norm'],
upsample_norm=opt_net['upsample_norm'])
elif which_model == 'modulate_denoise_resnet':
netG = arch.ModulateDenoiseResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
upscale=opt_net['scale'], norm_type=opt_net['norm_type'], mode=opt_net['mode'],
upsample_mode='pixelshuffle', ada_ksize=opt_net['ada_ksize'],
gate_conv_bias=opt_net['gate_conv_bias'])
elif which_model == 'noise_subnet':
netG = arch.NoiseSubNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
norm_type=opt_net['norm_type'], mode=opt_net['mode'])
elif which_model == 'cond_denoise_resnet':
netG = arch.CondDenoiseResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
upscale=opt_net['scale'], upsample_mode='pixelshuffle', ada_ksize=opt_net['ada_ksize'],
down_scale=opt_net['down_scale'], num_classes=opt_net['num_classes'],
norm_type=opt_net['norm_type'])
elif which_model == 'adabn_denoise_resnet':
netG = arch.AdaptiveDenoiseResNet(in_nc=opt_net['in_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
upscale=opt_net['scale'], down_scale=opt_net['down_scale'])
elif which_model == 'sft_arch': # SFT-GAN
netG = sft_arch.SFT_Net()
elif which_model == 'RRDB_net': # RRDB
netG = arch.RRDBNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'],
nb=opt_net['nb'], gc=opt_net['gc'], upscale=opt_net['scale'], norm_type=opt_net['norm_type'],
act_type='leakyrelu', mode=opt_net['mode'], upsample_mode='upconv')
else:
raise NotImplementedError('Generator model [{:s}] not recognized'.format(which_model))
if opt['init_type'] is not None:
init_weights(netG, init_type=opt['init_type'], scale=0.1)
if gpu_ids:
assert torch.cuda.is_available()
netG = nn.DataParallel(netG)
return netG
def define_sub(opt):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_sub']
which_model = opt_net['which_model_sub']
if which_model == 'noise_subnet':
subnet = arch.NoiseSubNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
norm_type=opt_net['norm_type'], mode=opt_net['mode'])
else:
raise NotImplementedError('subnet model [{:s}] not recognized'.format(which_model))
if gpu_ids:
assert torch.cuda.is_available()
subnet = nn.DataParallel(subnet)
return subnet
def define_sub2(opt):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_sub2']
which_model = opt_net['which_model_sub']
if which_model == 'blur_subnet':
subnet = arch.NoiseSubNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
norm_type=opt_net['norm_type'], mode=opt_net['mode'])
elif which_model == 'denoise_resnet':
subnet = arch.DenoiseResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], nb=opt_net['nb'],
upscale=opt_net['scale'], norm_type=opt_net['norm_type'], mode=opt_net['mode'],
upsample_mode='pixelshuffle', ada_ksize=opt_net['ada_ksize'],
down_scale=opt_net['down_scale'], fea_norm=opt_net['fea_norm'],
upsample_norm=opt_net['upsample_norm'])
else:
raise NotImplementedError('subnet model [{:s}] not recognized'.format(which_model))
if gpu_ids:
assert torch.cuda.is_available()
subnet = nn.DataParallel(subnet)
return subnet
# Discriminator
def define_D(opt):
gpu_ids = opt['gpu_ids']
opt_net = opt['network_D']
which_model = opt_net['which_model_D']
if which_model == 'discriminator_vgg_128':
netD = arch.Discriminator_VGG_128(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'dis_acd': # sft-gan, Auxiliary Classifier Discriminator
netD = sft_arch.ACD_VGG_BN_96()
elif which_model == 'discriminator_vgg_96':
netD = arch.Discriminator_VGG_96(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_192':
netD = arch.Discriminator_VGG_192(in_nc=opt_net['in_nc'], base_nf=opt_net['nf'], \
norm_type=opt_net['norm_type'], mode=opt_net['mode'], act_type=opt_net['act_type'])
elif which_model == 'discriminator_vgg_128_SN':
netD = arch.Discriminator_VGG_128_SN()
else:
raise NotImplementedError('Discriminator model [{:s}] not recognized'.format(which_model))
init_weights(netD, init_type='kaiming', scale=1)
if gpu_ids:
netD = nn.DataParallel(netD)
return netD
def define_F(opt, use_bn=False):
gpu_ids = opt['gpu_ids']
device = torch.device('cuda' if gpu_ids else 'cpu')
# pytorch pretrained VGG19-54, before ReLU.
if use_bn:
feature_layer = 49
else:
feature_layer = 34
netF = arch.VGGFeatureExtractor(feature_layer=feature_layer, use_bn=use_bn, \
use_input_norm=True, device=device)
# netF = arch.ResNet101FeatureExtractor(use_input_norm=True, device=device)
if gpu_ids:
netF = nn.DataParallel(netF)
netF.eval() # No need to train
return netF
网络结构
import math
import torch
import torch.nn as nn
import torchvision
import torch.nn.functional as F
from . import block as B
from . import spectral_norm as SN
from . import adaptive_norm as AN
####################
# Generator
####################
class SRCNN(nn.Module):
def __init__(self, in_nc, out_nc, nf, norm_type='batch', act_type='relu', mode='CNA', ada_ksize=None):
super(SRCNN, self).__init__()
fea_conv = B.conv_block(in_nc, nf, kernel_size=9, norm_type=norm_type, act_type=act_type, mode=mode
, ada_ksize=ada_ksize)
mapping_conv = B.conv_block(nf, nf // 2, kernel_size=1, norm_type=norm_type, act_type=act_type,
mode=mode, ada_ksize=ada_ksize)
HR_conv = B.conv_block(nf // 2, out_nc, kernel_size=5, norm_type=norm_type, act_type=None,
mode=mode, ada_ksize=ada_ksize)
self.model = B.sequential(fea_conv, mapping_conv, HR_conv)
def forward(self, x):
x = self.model(x)
return x
class ARCNN(nn.Module):
def __init__(self, in_nc, out_nc, nf, norm_type='batch', act_type='relu', mode='CNA', ada_ksize=None):
super(ARCNN, self).__init__()
fea_conv = B.conv_block(in_nc, nf, kernel_size=9, norm_type=norm_type, act_type=act_type, mode=mode
, ada_ksize=ada_ksize)
conv1 = B.conv_block(nf, nf // 2, kernel_size=7, norm_type=norm_type, act_type=act_type,
mode=mode, ada_ksize=ada_ksize)
conv2 = B.conv_block(nf // 2, nf // 4, kernel_size=1, norm_type=norm_type, act_type=act_type,
mode=mode, ada_ksize=ada_ksize)
HR_conv = B.conv_block(nf // 4, out_nc, kernel_size=5, norm_type=norm_type, act_type=None,
mode=mode, ada_ksize=ada_ksize)
self.model = B.sequential(fea_conv, conv1, conv2, HR_conv)
def forward(self, x):
x = self.model(x)
return x
class SRResNet(nn.Module):
def __init__(self, in_nc, out_nc, nf, nb, upscale=4, norm_type='batch', act_type='relu', \
mode='NAC', res_scale=1, upsample_mode='upconv'):
super(SRResNet, self).__init__()
n_upscale = int(math.log(upscale, 2))
if upscale == 3:
n_upscale = 1
fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None)
resnet_blocks = [B.ResNetBlock(nf, nf, nf, norm_type=norm_type, act_type=act_type,\
mode=mode, res_scale=res_scale) for _ in range(nb)]
LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode)
if upsample_mode == 'upconv':
upsample_block = B.upconv_blcok
elif upsample_mode == 'pixelshuffle':
upsample_block = B.pixelshuffle_block
else:
raise NotImplementedError('upsample mode [{:s}] is not found'.format(upsample_mode))
if upscale == 3:
upsampler = upsample_block(nf, nf, 3, act_type=act_type)
else:
upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*resnet_blocks, LR_conv)),\
*upsampler, HR_conv0, HR_conv1)
def forward(self, x):
x = self.model(x)
return x
class ModulateSRResNet(nn.Module):
def __init__(self, in_nc, out_nc, nf, nb, upscale=4, norm_type='sft', act_type='relu',
mode='CNA', res_scale=1, upsample_mode='upconv', gate_conv_bias=True, ada_ksize=None):
super(ModulateSRResNet, self).__init__()
n_upscale = int(math.log(upscale, 2))
if upscale == 3:
n_upscale = 1
self.fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None, stride=1)
resnet_blocks = [B.TwoStreamSRResNet(nf, nf, nf, norm_type=norm_type, act_type=act_type,
mode=mode, res_scale=res_scale, gate_conv_bias=gate_conv_bias,
ada_ksize=ada_ksize, input_dim=in_nc) for _ in range(nb)]
self.LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=None, mode=mode)
if norm_type == 'sft':
self.LR_norm = AN.GateNonLinearLayer(in_nc, conv_bias=gate_conv_bias)
elif norm_type == 'sft_conv':
self.LR_norm = AN.MetaLayer(in_nc, conv_bias=gate_conv_bias, kernel_size=ada_ksize)
if upsample_mode == 'upconv':
upsample_block = B.upconv_blcok
elif upsample_mode == 'pixelshuffle':
upsample_block = B.pixelshuffle_block
else:
raise NotImplementedError('upsample mode [%s] is not found' % upsample_mode)
if upscale == 3:
upsampler = upsample_block(nf, nf, 3, act_type=act_type)
else:
upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
self.norm_branch = B.sequential(*resnet_blocks)
self.HR_branch = B.sequential(*upsampler, HR_conv0, HR_conv1)
def forward(self, x):
fea = self.fea_conv(x[0])
fea_res_block, _ = self.norm_branch((fea, x[1]))
fea_LR = self.LR_conv(fea_res_block)
res = self.LR_norm((fea_LR, x[1]))
out = self.HR_branch(fea+res)
return out
class DenoiseResNet(nn.Module):
"""
jingwen's addition
denoise Resnet
"""
def __init__(self, in_nc, out_nc, nf, nb, upscale=1, norm_type='batch', act_type='relu',
mode='CNA', res_scale=1, upsample_mode='upconv', ada_ksize=None, down_scale=2,
fea_norm=None, upsample_norm=None):
super(DenoiseResNet, self).__init__()
n_upscale = int(math.log(down_scale, 2))
if down_scale == 3:
n_upscale = 1
fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=fea_norm, act_type=None, stride=down_scale,
ada_ksize=ada_ksize)
resnet_blocks = [B.ResNetBlock(nf, nf, nf, norm_type=norm_type, act_type=act_type,
mode=mode, res_scale=res_scale, ada_ksize=ada_ksize) for _ in range(nb)]
LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=None, mode=mode
, ada_ksize=ada_ksize)
# LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=None, mode=mode
# , ada_ksize=ada_ksize)
if upsample_mode == 'upconv':
upsample_block = B.upconv_blcok
elif upsample_mode == 'pixelshuffle':
upsample_block = B.pixelshuffle_block
else:
raise NotImplementedError('upsample mode [%s] is not found' % upsample_mode)
if down_scale == 3:
upsampler = upsample_block(nf, nf, 3, act_type=act_type, norm_type=upsample_norm, ada_ksize=ada_ksize)
else:
upsampler = [upsample_block(nf, nf, act_type=act_type, norm_type=upsample_norm, ada_ksize=ada_ksize) for _ in range(n_upscale)]
HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=upsample_norm, act_type=act_type, ada_ksize=ada_ksize)
HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=upsample_norm, act_type=None, ada_ksize=ada_ksize)
self.model = B.sequential(fea_conv, B.ShortcutBlock(B.sequential(*resnet_blocks, LR_conv)),
*upsampler, HR_conv0, HR_conv1)
def forward(self, x):
x = self.model(x)
return x
class ModulateDenoiseResNet(nn.Module):
def __init__(self, in_nc, out_nc, nf, nb, upscale=1, norm_type='sft', act_type='relu',
mode='CNA', res_scale=1, upsample_mode='upconv', gate_conv_bias=True, ada_ksize=None):
super(ModulateDenoiseResNet, self).__init__()
self.fea_conv = B.conv_block(in_nc, nf, kernel_size=3, norm_type=None, act_type=None, stride=2)
resnet_blocks = [B.TwoStreamSRResNet(nf, nf, nf, norm_type=norm_type, act_type=act_type,
mode=mode, res_scale=res_scale, gate_conv_bias=gate_conv_bias,
ada_ksize=ada_ksize, input_dim=in_nc) for _ in range(nb)]
LR_conv = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=None, mode=mode)
if norm_type == 'sft':
LR_norm = AN.GateNonLinearLayer(in_nc, conv_bias=gate_conv_bias)
elif norm_type == 'sft_conv':
LR_norm = AN.MetaLayer(in_nc, conv_bias=gate_conv_bias, kernel_size=ada_ksize)
if upsample_mode == 'upconv':
upsample_block = B.upconv_blcok
elif upsample_mode == 'pixelshuffle':
upsample_block = B.pixelshuffle_block
else:
raise NotImplementedError('upsample mode [%s] is not found' % upsample_mode)
upsampler = upsample_block(nf, nf, act_type=act_type)
HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
self.norm_branch = B.sequential(*resnet_blocks)
self.LR_conv = LR_conv
self.LR_norm = LR_norm
self.HR_branch = B.sequential(upsampler, HR_conv0, HR_conv1)
def forward(self, x):
fea = self.fea_conv(x[0])
fea_res_block, _ = self.norm_branch((fea, x[1]))
fea_LR = self.LR_conv(fea_res_block)
res = self.LR_norm((fea_LR, x[1]))
out = self.HR_branch(fea+res)
return out
class NoiseSubNet(nn.Module):
def __init__(self, in_nc, out_nc, nf, nb, norm_type='batch', act_type='relu', mode='CNA'):
super(NoiseSubNet, self).__init__()
degration_block = [B.conv_block(in_nc, nf, kernel_size=3, norm_type=norm_type, act_type=act_type, mode=mode)]
degration_block.extend([B.conv_block(nf, nf, kernel_size=3, norm_type=norm_type, act_type=act_type, mode=mode)
for _ in range(15)])
degration_block.append(B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None, mode=mode))
self.degration_block = B.sequential(*degration_block)
def forward(self, x):
deg_estimate = self.degration_block(x)
return deg_estimate
class CondDenoiseResNet(nn.Module):
"""
jingwen's addition
denoise Resnet
"""
def __init__(self, in_nc, out_nc, nf, nb, upscale=1, res_scale=1, down_scale=2, num_classes=1, ada_ksize=None
,upsample_mode='upconv', act_type='relu', norm_type='cond_adaptive_conv_res'):
super(CondDenoiseResNet, self).__init__()
n_upscale = int(math.log(down_scale, 2))
if down_scale == 3:
n_upscale = 1
self.fea_conv = nn.Conv2d(in_nc, nf, kernel_size=3, stride=down_scale, padding=1)
resnet_blocks = [B.CondResNetBlock(nf, nf, nf, num_classes=num_classes, ada_ksize=ada_ksize,
norm_type=norm_type, act_type=act_type) for _ in range(nb)]
self.resnet_blocks = B.sequential(*resnet_blocks)
self.LR_conv = nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1)
if norm_type == 'cond_adaptive_conv_res':
self.cond_adaptive = AN.CondAdaptiveConvResNorm(nf, num_classes=num_classes)
elif norm_type == "interp_adaptive_conv_res":
self.cond_adaptive = AN.InterpAdaptiveResNorm(nf, ada_ksize)
elif norm_type == "cond_instance":
self.cond_adaptive = AN.CondInstanceNorm2d(nf, num_classes=num_classes)
elif norm_type == "cond_transform_res":
self.cond_adaptive = AN.CondResTransformer(nf, ada_ksize, num_classes=num_classes)
if upsample_mode == 'upconv':
upsample_block = B.upconv_blcok
elif upsample_mode == 'pixelshuffle':
upsample_block = B.pixelshuffle_block
else:
raise NotImplementedError('upsample mode [%s] is not found' % upsample_mode)
if down_scale == 3:
upsampler = upsample_block(nf, nf, 3, act_type=act_type)
else:
upsampler = [upsample_block(nf, nf, act_type=act_type) for _ in range(n_upscale)]
HR_conv0 = B.conv_block(nf, nf, kernel_size=3, norm_type=None, act_type=act_type)
HR_conv1 = B.conv_block(nf, out_nc, kernel_size=3, norm_type=None, act_type=None)
self.upsample = B.sequential(*upsampler, HR_conv0, HR_conv1)
def forward(self, x, y):
# the first feature extraction
fea = self.fea_conv(x)
fea1, _ = self.resnet_blocks((fea, y))
fea2 = self.LR_conv(fea1)
fea3 = self.cond_adaptive(fea2, y)
# res
out = self.upsample(fea3 + fea)
return out
class AdaptiveDenoiseResNet(nn.Module):
"""
jingwen's addition
adabn
"""
def __init__(self, in_nc, nf, nb, upscale=1, res_scale=1, down_scale=2):
super(AdaptiveDenoiseResNet, self).__init__()
self.fea_conv = nn.Conv2d(in_nc, nf, kernel_size=3, stride=down_scale, padding=1)
resnet_blocks = [B.AdaptiveResNetBlock(nf, nf, nf, res_scale=res_scale) for _ in range(nb)]
self.resnet_blocks = B.sequential(*resnet_blocks)
self.LR_conv = nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1)
self.batch_norm = nn.BatchNorm2d(nf, affine=True, track_running_stats=True, momentum=0)
def forward(self, x):
fea_list = [self.fea_conv(data.unsqueeze_(0)) for data in x]
fea_resblock_list = self.resnet_blocks(fea_list)
fea_LR_list = [self.LR_conv(fea) for fea in fea_resblock_list]
fea_mean, fea_var = B.computing_mean_variance(fea_LR_list)
batch_norm_dict = self.batch_norm.state_dict()
batch_norm_dict['running_mean'] = fea_mean
batch_norm_dict['running_var'] = fea_var
self.batch_norm.load_state_dict(batch_norm_dict)
return None
experiment
