From 0d070b47a7001c83de90d5f2698926fb6119f940 Mon Sep 17 00:00:00 2001 From: James Betker Date: Mon, 3 Aug 2020 10:25:37 -0600 Subject: [PATCH] Add simplified SPSR architecture Basically just cleaning up the code, removing some bad conventions, and reducing complexity somewhat so that I can play around with this arch a bit more easily. --- codes/models/SRGAN_model.py | 11 ++- codes/models/archs/SPSR_arch.py | 116 +++++++++++++++++++++++++++++++- codes/models/archs/SPSR_util.py | 4 +- codes/models/archs/arch_util.py | 29 ++++++++ codes/models/networks.py | 3 + 5 files changed, 153 insertions(+), 10 deletions(-) diff --git a/codes/models/SRGAN_model.py b/codes/models/SRGAN_model.py index 14bd293c..93798d43 100644 --- a/codes/models/SRGAN_model.py +++ b/codes/models/SRGAN_model.py @@ -394,9 +394,8 @@ class SRGANModel(BaseModel): fea_GenOut = fake_GenOut using_gan_img = False # Get image gradients for later use. - fake_H_grad = self.get_grad(fake_GenOut) - var_H_grad = self.get_grad(var_H) - var_ref_grad = self.get_grad(var_ref) + fake_H_grad = self.get_grad_nopadding(fake_GenOut) + var_ref_grad = self.get_grad_nopadding(var_ref) var_H_grad_nopadding = self.get_grad_nopadding(var_H) self.spsr_grad_GenOut.append(grad_LR) else: @@ -426,7 +425,7 @@ class SRGANModel(BaseModel): l_g_pix_log = l_g_pix / self.l_pix_w l_g_total += l_g_pix if self.spsr_enabled and self.cri_pix_grad: # gradient pixel loss - l_g_pix_grad = self.l_pix_grad_w * self.cri_pix_grad(fake_H_grad, var_H_grad) + l_g_pix_grad = self.l_pix_grad_w * self.cri_pix_grad(fake_H_grad, var_H_grad_nopadding) l_g_total += l_g_pix_grad if self.spsr_enabled and self.cri_pix_branch: # branch pixel loss l_g_pix_grad_branch = self.l_pix_branch_w * self.cri_pix_branch(fake_H_branch, @@ -660,8 +659,8 @@ class SRGANModel(BaseModel): self.optimizer_D_grad.zero_grad() for var_ref, fake_H in zip(self.var_ref, self.fake_H): - fake_H_grad = self.get_grad(fake_H) - var_ref_grad = self.get_grad(var_ref) + fake_H_grad = self.get_grad_nopadding(fake_H) + var_ref_grad = self.get_grad_nopadding(var_ref) pred_d_real_grad = self.netD_grad(var_ref_grad) pred_d_fake_grad = self.netD_grad(fake_H_grad.detach()) # detach to avoid BP to G if self.opt['train']['gan_type'] == 'gan': diff --git a/codes/models/archs/SPSR_arch.py b/codes/models/archs/SPSR_arch.py index 90982e9f..1757d8c9 100644 --- a/codes/models/archs/SPSR_arch.py +++ b/codes/models/archs/SPSR_arch.py @@ -4,6 +4,7 @@ import torch.nn as nn import torch.nn.functional as F from models.archs import SPSR_util as B from .RRDBNet_arch import RRDB +from models.archs.arch_util import ConvGnLelu, ExpansionBlock, UpconvBlock class ImageGradient(nn.Module): @@ -91,7 +92,7 @@ class SPSRNet(nn.Module): 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 + upsample_block = B.upconv_block elif upsample_mode == 'pixelshuffle': upsample_block = B.pixelshuffle_block else: @@ -129,7 +130,7 @@ class SPSRNet(nn.Module): self.b_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 + upsample_block = B.upconv_block elif upsample_mode == 'pixelshuffle': upsample_block = B.pixelshuffle_block else: @@ -224,3 +225,114 @@ class SPSRNet(nn.Module): ######### return x_out_branch, x_out, x_grad + +class SPSRNetSimplified(nn.Module): + def __init__(self, in_nc, out_nc, nf, nb, upscale=4): + super(SPSRNetSimplified, self).__init__() + n_upscale = int(math.log(upscale, 2)) + + # Feature branch + self.model_fea_conv = ConvGnLelu(in_nc, nf, kernel_size=3, norm=False, activation=False) + self.model_shortcut_blk = nn.Sequential(*[RRDB(nf, gc=32) for _ in range(nb)]) + self.feature_lr_conv = ConvGnLelu(nf, nf, kernel_size=3, norm=False, activation=False) + self.model_upsampler = nn.Sequential(*[UpconvBlock(nf, nf, block=ConvGnLelu, norm=False, activation=False, bias=False) for _ in range(n_upscale)]) + self.feature_hr_conv1 = ConvGnLelu(nf, nf, kernel_size=3, norm=False, activation=True, bias=False) + self.feature_hr_conv2 = ConvGnLelu(nf, nf, kernel_size=3, norm=False, activation=False, bias=False) + + # Grad branch + self.get_g_nopadding = ImageGradientNoPadding() + self.b_fea_conv = ConvGnLelu(in_nc, nf, kernel_size=3, norm=False, activation=False, bias=False) + self.b_concat_decimate_1 = ConvGnLelu(2 * nf, nf, kernel_size=1, norm=False, activation=False, bias=False) + self.b_proc_block_1 = RRDB(nf, gc=32) + self.b_concat_decimate_2 = ConvGnLelu(2 * nf, nf, kernel_size=1, norm=False, activation=False, bias=False) + self.b_proc_block_2 = RRDB(nf, gc=32) + self.b_concat_decimate_3 = ConvGnLelu(2 * nf, nf, kernel_size=1, norm=False, activation=False, bias=False) + self.b_proc_block_3 = RRDB(nf, gc=32) + self.b_concat_decimate_4 = ConvGnLelu(2 * nf, nf, kernel_size=1, norm=False, activation=False, bias=False) + self.b_proc_block_4 = RRDB(nf, gc=32) + # Upsampling + self.grad_lr_conv = ConvGnLelu(nf, nf, kernel_size=3, norm=False, activation=True, bias=False) + b_upsampler = nn.Sequential(*[UpconvBlock(nf, nf, block=ConvGnLelu, norm=False, activation=False, bias=False) for _ in range(n_upscale)]) + grad_hr_conv1 = ConvGnLelu(nf, nf, kernel_size=3, norm=False, activation=True, bias=False) + grad_hr_conv2 = ConvGnLelu(nf, nf, kernel_size=3, norm=False, activation=False, bias=False) + self.branch_upsample = B.sequential(*b_upsampler, grad_hr_conv1, grad_hr_conv2) + # Conv used to output grad branch shortcut. + self.grad_branch_output_conv = ConvGnLelu(nf, out_nc, kernel_size=1, norm=False, activation=False, bias=False) + + # Conjoin branch. + # Note: "_branch_pretrain" is a special tag used to denote parameters that get pretrained before the rest. + self._branch_pretrain_concat = ConvGnLelu(nf * 2, nf, kernel_size=1, norm=False, activation=False, bias=False) + self._branch_pretrain_block = RRDB(nf * 2, gc=32) + self._branch_pretrain_HR_conv0 = ConvGnLelu(nf, nf, kernel_size=3, norm=False, activation=True, bias=False) + self._branch_pretrain_HR_conv1 = ConvGnLelu(nf, out_nc, kernel_size=3, norm=False, activation=False, bias=False) + + def forward(self, x): + + x_grad = self.get_g_nopadding(x) + x = self.model_fea_conv(x) + + x_ori = x + for i in range(5): + x = self.model_shortcut_blk[i](x) + x_fea1 = x + + for i in range(5): + x = self.model_shortcut_blk[i + 5](x) + x_fea2 = x + + for i in range(5): + x = self.model_shortcut_blk[i + 10](x) + x_fea3 = x + + for i in range(5): + x = self.model_shortcut_blk[i + 15](x) + x_fea4 = x + + x = self.model_shortcut_blk[20:](x) + x = self.feature_lr_conv(x) + + # short cut + x = x_ori + x + x = self.model_upsampler(x) + x = self.feature_hr_conv1(x) + x = self.feature_hr_conv2(x) + + x_b_fea = self.b_fea_conv(x_grad) + x_cat_1 = torch.cat([x_b_fea, x_fea1], dim=1) + + x_cat_1 = self.b_concat_decimate_1(x_cat_1) + x_cat_1 = self.b_proc_block_1(x_cat_1) + + x_cat_2 = torch.cat([x_cat_1, x_fea2], dim=1) + + x_cat_2 = self.b_concat_decimate_2(x_cat_2) + x_cat_2 = self.b_proc_block_2(x_cat_2) + + x_cat_3 = torch.cat([x_cat_2, x_fea3], dim=1) + + x_cat_3 = self.b_concat_decimate_3(x_cat_3) + x_cat_3 = self.b_proc_block_3(x_cat_3) + + x_cat_4 = torch.cat([x_cat_3, x_fea4], dim=1) + + x_cat_4 = self.b_concat_decimate_4(x_cat_4) + x_cat_4 = self.b_proc_block_4(x_cat_4) + + x_cat_4 = self.grad_lr_conv(x_cat_4) + + # short cut + x_cat_4 = x_cat_4 + x_b_fea + x_branch = self.branch_upsample(x_cat_4) + x_out_branch = self.grad_branch_output_conv(x_branch) + + ######## + x_branch_d = x_branch + x__branch_pretrain_cat = torch.cat([x_branch_d, x], dim=1) + x__branch_pretrain_cat = self._branch_pretrain_block(x__branch_pretrain_cat) + x_out = self._branch_pretrain_concat(x__branch_pretrain_cat) + x_out = self._branch_pretrain_HR_conv0(x_out) + x_out = self._branch_pretrain_HR_conv1(x_out) + + ######### + return x_out_branch, x_out, x_grad + diff --git a/codes/models/archs/SPSR_util.py b/codes/models/archs/SPSR_util.py index 0485f5ae..df40325f 100644 --- a/codes/models/archs/SPSR_util.py +++ b/codes/models/archs/SPSR_util.py @@ -153,8 +153,8 @@ def pixelshuffle_block(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, return sequential(conv, pixel_shuffle, n, a) -def upconv_blcok(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, bias=True, \ - pad_type='zero', norm_type=None, act_type='relu', mode='nearest'): +def upconv_block(in_nc, out_nc, upscale_factor=2, kernel_size=3, stride=1, bias=True, \ + pad_type='zero', norm_type=None, act_type='relu', mode='nearest'): # Up conv # described in https://distill.pub/2016/deconv-checkerboard/ upsample = nn.Upsample(scale_factor=upscale_factor, mode=mode) diff --git a/codes/models/archs/arch_util.py b/codes/models/archs/arch_util.py index fba27c98..1c1143ad 100644 --- a/codes/models/archs/arch_util.py +++ b/codes/models/archs/arch_util.py @@ -388,4 +388,33 @@ class ExpansionBlock(nn.Module): x = self.decimate(x) p = self.process_passthrough(passthrough) x = self.conjoin(torch.cat([x, p], dim=1)) + return self.process(x) + + +# Similar to ExpansionBlock but does not upsample. +class ConjoinBlock(nn.Module): + def __init__(self, filters_in, filters_out=None, block=ConvGnSilu, norm=True): + super(ConjoinBlock, self).__init__() + if filters_out is None: + filters_out = filters_in + self.decimate = block(filters_in*2, filters_out, kernel_size=1, bias=False, activation=False, norm=norm) + self.process = block(filters_out, filters_out, kernel_size=3, bias=False, activation=True, norm=norm) + + # input is the feature signal with shape (b, f, w, h) + # passthrough is the structure signal with shape (b, f/2, w*2, h*2) + # output is conjoined upsample with shape (b, f/2, w*2, h*2) + def forward(self, input, passthrough): + x = torch.cat([input, passthrough], dim=1) + x = self.decimate(x) + return self.process(x) + + +# Basic convolutional upsampling block that uses interpolate. +class UpconvBlock(nn.Module): + def __init__(self, filters_in, filters_out=None, block=ConvGnSilu, norm=True, activation=True, bias=False): + super(UpconvBlock, self).__init__() + self.process = block(filters_out, filters_out, kernel_size=3, bias=bias, activation=activation, norm=norm) + + def forward(self, x): + x = F.interpolate(x, scale_factor=2, mode="nearest") return self.process(x) \ No newline at end of file diff --git a/codes/models/networks.py b/codes/models/networks.py index 07aaa4c1..5fccc92b 100644 --- a/codes/models/networks.py +++ b/codes/models/networks.py @@ -105,6 +105,9 @@ def define_G(opt, net_key='network_G'): act_type='leakyrelu', mode=opt_net['mode'], upsample_mode='upconv') if opt['is_train']: arch_util.initialize_weights(netG, scale=.1) + elif which_model == 'spsr_net_improved': + netG = spsr.SPSRNetSimplified(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], nf=opt_net['nf'], + nb=opt_net['nb'], upscale=opt_net['scale']) # image corruption elif which_model == 'HighToLowResNet':