diff --git a/codes/models/SRGAN_model.py b/codes/models/SRGAN_model.py index 1d13bc95..902fa226 100644 --- a/codes/models/SRGAN_model.py +++ b/codes/models/SRGAN_model.py @@ -494,10 +494,11 @@ class SRGANModel(BaseModel): if self.spsr_enabled and self.cri_grad_gan: if self.opt['train']['gan_type'] == 'crossgan': - pred_g_fake_grad = self.netD(fake_H_grad, var_L) + pred_g_fake_grad = self.netD_grad(fake_H_grad, var_L) + pred_g_fake_grad_branch = self.netD_grad(fake_H_branch, var_L) else: - pred_g_fake_grad = self.netD(fake_H_grad) - pred_g_fake_grad_branch = self.netD_grad(fake_H_branch) + pred_g_fake_grad = self.netD_grad(fake_H_grad) + pred_g_fake_grad_branch = self.netD_grad(fake_H_branch) if self.opt['train']['gan_type'] in ['gan', 'pixgan', 'pixgan_fea', 'crossgan']: l_g_gan_grad = self.l_gan_grad_w * self.cri_grad_gan(pred_g_fake_grad, True) l_g_gan_grad_branch = self.l_gan_grad_w * self.cri_grad_gan(pred_g_fake_grad_branch, True) @@ -685,19 +686,23 @@ class SRGANModel(BaseModel): for p in self.netD_grad.parameters(): p.requires_grad = True self.optimizer_D_grad.zero_grad() - for var_ref, fake_H, fake_H_grad_branch in zip(var_ref_skips, self.fake_H, self.spsr_grad_GenOut): + for var_L, var_ref, fake_H, fake_H_grad_branch in zip(self.var_L, var_ref_skips, self.fake_H, self.spsr_grad_GenOut): fake_H_grad = self.get_grad_nopadding(fake_H).detach() 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) # Tensor already detached above. - # var_ref and fake_H already has noise added to it. We **must** add noise to fake_H_grad_branch too. fake_H_grad_branch = fake_H_grad_branch.detach() + noise - pred_d_fake_grad_branch = self.netD_grad(fake_H_grad_branch) - if self.opt['train']['gan_type'] == 'gan': + if self.opt['train']['gan_type'] == 'crossgan': + pred_d_real_grad = self.netD_grad(var_ref_grad, var_L) + pred_d_fake_grad = self.netD_grad(fake_H_grad, var_L) # Tensor already detached above. + # var_ref and fake_H already has noise added to it. We **must** add noise to fake_H_grad_branch too. + pred_d_fake_grad_branch = self.netD_grad(fake_H_grad_branch, var_L) + else: + pred_d_real_grad = self.netD_grad(var_ref_grad) + pred_d_fake_grad = self.netD_grad(fake_H_grad) # Tensor already detached above. + # var_ref and fake_H already has noise added to it. We **must** add noise to fake_H_grad_branch too. + pred_d_fake_grad_branch = self.netD_grad(fake_H_grad_branch) + if self.opt['train']['gan_type'] == 'gan' or self.opt['train']['gan_type'] == 'crossgan': l_d_real_grad = self.cri_gan(pred_d_real_grad, True) l_d_fake_grad = (self.cri_gan(pred_d_fake_grad, False) + self.cri_gan(pred_d_fake_grad_branch, False)) / 2 - elif self.opt['train']['gan_type'] == 'crossgan': - assert False elif self.opt['train']['gan_type'] == 'pixgan': real = torch.ones_like(pred_d_real_grad) fake = torch.zeros_like(pred_d_fake_grad) diff --git a/codes/models/archs/SPSR_arch.py b/codes/models/archs/SPSR_arch.py index f23ee624..dec5d1ea 100644 --- a/codes/models/archs/SPSR_arch.py +++ b/codes/models/archs/SPSR_arch.py @@ -670,3 +670,125 @@ class SwitchedSpsrLr(nn.Module): val["switch_%i_specificity" % (i,)] = means[i] val["switch_%i_histogram" % (i,)] = hists[i] return val + + +class SwitchedSpsrLr2(nn.Module): + def __init__(self, in_nc, out_nc, nf, upscale=4): + super(SwitchedSpsrLr2, self).__init__() + n_upscale = int(math.log(upscale, 2)) + + # switch options + transformation_filters = nf + switch_filters = nf + switch_reductions = 3 + switch_processing_layers = 2 + self.transformation_counts = 8 + multiplx_fn = functools.partial(ConvBasisMultiplexer, transformation_filters, switch_filters, switch_reductions, + switch_processing_layers, self.transformation_counts, use_exp2=True) + pretransform_fn = functools.partial(ConvGnLelu, transformation_filters, transformation_filters, norm=False, bias=False, weight_init_factor=.1) + transform_fn = functools.partial(MultiConvBlock, transformation_filters, int(transformation_filters * 1.5), + transformation_filters, kernel_size=3, depth=3, + weight_init_factor=.1) + + # Feature branch + self.model_fea_conv = ConvGnLelu(in_nc, nf, kernel_size=3, norm=False, activation=False) + self.sw1 = ConfigurableSwitchComputer(transformation_filters, multiplx_fn, + pre_transform_block=pretransform_fn, transform_block=transform_fn, + attention_norm=True, + transform_count=self.transformation_counts, init_temp=10, + add_scalable_noise_to_transforms=True) + self.sw2 = ConfigurableSwitchComputer(transformation_filters, multiplx_fn, + pre_transform_block=pretransform_fn, transform_block=transform_fn, + attention_norm=True, + transform_count=self.transformation_counts, init_temp=10, + add_scalable_noise_to_transforms=True) + self.feature_lr_conv = ConvGnLelu(nf, nf, kernel_size=3, norm=True, activation=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) + mplex_grad = functools.partial(ConvBasisMultiplexer, nf * 2, nf * 2, switch_reductions, + switch_processing_layers, self.transformation_counts // 2, use_exp2=True) + self.sw_grad = ConfigurableSwitchComputer(transformation_filters, mplex_grad, + pre_transform_block=pretransform_fn, transform_block=transform_fn, + attention_norm=True, + transform_count=self.transformation_counts // 2, init_temp=10, + add_scalable_noise_to_transforms=True) + # Upsampling + self.grad_lr_conv = ConvGnLelu(nf, nf, kernel_size=3, norm=True, activation=True, bias=False) + self.grad_hr_conv = ConvGnLelu(nf, nf, kernel_size=3, norm=False, activation=False, bias=False) + # 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. + transform_fn_cat = functools.partial(MultiConvBlock, transformation_filters * 2, int(transformation_filters * 1.5), + transformation_filters, kernel_size=3, depth=4, + weight_init_factor=.1) + pretransform_fn_cat = functools.partial(ConvGnLelu, transformation_filters * 2, transformation_filters * 2, norm=False, bias=False, weight_init_factor=.1) + self._branch_pretrain_sw = ConfigurableSwitchComputer(transformation_filters, multiplx_fn, + pre_transform_block=pretransform_fn_cat, transform_block=transform_fn_cat, + attention_norm=True, + transform_count=self.transformation_counts, init_temp=10, + add_scalable_noise_to_transforms=True) + self.upsample = nn.Sequential(*[UpconvBlock(nf, nf, block=ConvGnLelu, norm=False, activation=False, bias=False) for _ in range(n_upscale)]) + self.upsample_grad = nn.Sequential(*[UpconvBlock(nf, nf, block=ConvGnLelu, norm=False, activation=False, bias=False) for _ in range(n_upscale)]) + self.final_lr_conv = ConvGnLelu(nf, nf, kernel_size=3, norm=True, activation=False) + self.final_hr_conv1 = ConvGnLelu(nf, nf, kernel_size=3, norm=True, activation=True, bias=False) + self.final_hr_conv2 = ConvGnLelu(nf, out_nc, kernel_size=3, norm=False, activation=False, bias=False) + self.switches = [self.sw1, self.sw2, self.sw_grad, self._branch_pretrain_sw] + self.attentions = None + self.init_temperature = 10 + self.final_temperature_step = 10000 + + def forward(self, x): + x_grad = self.get_g_nopadding(x) + x = self.model_fea_conv(x) + + x1, a1 = self.sw1(x, True) + x2, a2 = self.sw2(x1, True) + x_fea = self.feature_lr_conv(x2) + x_fea = self.feature_hr_conv2(x_fea) + + x_b_fea = self.b_fea_conv(x_grad) + x_grad, a3 = self.sw_grad(x_b_fea, att_in=torch.cat([x1, x_b_fea], dim=1), output_attention_weights=True) + x_grad = self.grad_lr_conv(x_grad) + x_grad = self.grad_hr_conv(x_grad) + x_out_branch = self.upsample_grad(x_grad) + x_out_branch = self.grad_branch_output_conv(x_out_branch) + + x__branch_pretrain_cat = torch.cat([x_grad, x_fea], dim=1) + x__branch_pretrain_cat, a4 = self._branch_pretrain_sw(x__branch_pretrain_cat, att_in=x_fea, identity=x_fea, output_attention_weights=True) + x_out = self.final_lr_conv(x__branch_pretrain_cat) + x_out = self.upsample(x_out) + x_out = self.final_hr_conv1(x_out) + x_out = self.final_hr_conv2(x_out) + + self.attentions = [a1, a2, a3, a4] + + return x_out_branch, x_out, x_grad + + def set_temperature(self, temp): + [sw.set_temperature(temp) for sw in self.switches] + + def update_for_step(self, step, experiments_path='.'): + if self.attentions: + temp = max(1, 1 + self.init_temperature * + (self.final_temperature_step - step) / self.final_temperature_step) + self.set_temperature(temp) + if step % 200 == 0: + output_path = os.path.join(experiments_path, "attention_maps", "a%i") + prefix = "attention_map_%i_%%i.png" % (step,) + [save_attention_to_image_rgb(output_path % (i,), self.attentions[i], self.transformation_counts, prefix, step) for i in range(len(self.attentions))] + + def get_debug_values(self, step): + temp = self.switches[0].switch.temperature + mean_hists = [compute_attention_specificity(att, 2) for att in self.attentions] + means = [i[0] for i in mean_hists] + hists = [i[1].clone().detach().cpu().flatten() for i in mean_hists] + val = {"switch_temperature": temp} + for i in range(len(means)): + val["switch_%i_specificity" % (i,)] = means[i] + val["switch_%i_histogram" % (i,)] = hists[i] + return val diff --git a/codes/models/archs/SwitchedResidualGenerator_arch.py b/codes/models/archs/SwitchedResidualGenerator_arch.py index 3548211e..1fe3ba3a 100644 --- a/codes/models/archs/SwitchedResidualGenerator_arch.py +++ b/codes/models/archs/SwitchedResidualGenerator_arch.py @@ -4,7 +4,7 @@ from switched_conv import BareConvSwitch, compute_attention_specificity, Attenti import torch.nn.functional as F import functools from collections import OrderedDict -from models.archs.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock +from models.archs.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock, ExpansionBlock2 from models.archs.RRDBNet_arch import ResidualDenseBlock_5C, RRDB from models.archs.spinenet_arch import SpineNet from switched_conv_util import save_attention_to_image_rgb @@ -47,13 +47,16 @@ class HalvingProcessingBlock(nn.Module): # This is a classic u-net architecture with the goal of assigning each individual pixel an individual transform # switching set. class ConvBasisMultiplexer(nn.Module): - def __init__(self, input_channels, base_filters, reductions, processing_depth, multiplexer_channels, use_gn=True): + def __init__(self, input_channels, base_filters, reductions, processing_depth, multiplexer_channels, use_gn=True, use_exp2=False): super(ConvBasisMultiplexer, self).__init__() self.filter_conv = ConvGnSilu(input_channels, base_filters, bias=True) self.reduction_blocks = nn.ModuleList([HalvingProcessingBlock(base_filters * 2 ** i) for i in range(reductions)]) reduction_filters = base_filters * 2 ** reductions self.processing_blocks = nn.Sequential(OrderedDict([('block%i' % (i,), ConvGnSilu(reduction_filters, reduction_filters, bias=False)) for i in range(processing_depth)])) - self.expansion_blocks = nn.ModuleList([ExpansionBlock(reduction_filters // (2 ** i)) for i in range(reductions)]) + if use_exp2: + self.expansion_blocks = nn.ModuleList([ExpansionBlock2(reduction_filters // (2 ** i)) for i in range(reductions)]) + else: + self.expansion_blocks = nn.ModuleList([ExpansionBlock(reduction_filters // (2 ** i)) for i in range(reductions)]) gap = base_filters - multiplexer_channels cbl1_out = ((base_filters - (gap // 2)) // 4) * 4 # Must be multiples of 4 to use with group norm. diff --git a/codes/models/archs/arch_util.py b/codes/models/archs/arch_util.py index 1c1143ad..5af4e6ea 100644 --- a/codes/models/archs/arch_util.py +++ b/codes/models/archs/arch_util.py @@ -391,6 +391,30 @@ class ExpansionBlock(nn.Module): return self.process(x) +# Block that upsamples 2x and reduces incoming filters by 2x. It preserves structure by taking a passthrough feed +# along with the feature representation. +# Differs from ExpansionBlock because it performs all processing in 2xfilter space and decimates at the last step. +class ExpansionBlock2(nn.Module): + def __init__(self, filters_in, filters_out=None, block=ConvGnSilu): + super(ExpansionBlock2, self).__init__() + if filters_out is None: + filters_out = filters_in // 2 + self.decimate = block(filters_in, filters_out, kernel_size=1, bias=False, activation=False, norm=True) + self.process_passthrough = block(filters_out, filters_out, kernel_size=3, bias=True, activation=False, norm=True) + self.conjoin = block(filters_out*2, filters_out*2, kernel_size=3, bias=False, activation=True, norm=False) + self.reduce = block(filters_out*2, filters_out, kernel_size=3, bias=False, activation=True, norm=True) + + # 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 = F.interpolate(input, scale_factor=2, mode="nearest") + x = self.decimate(x) + p = self.process_passthrough(passthrough) + x = self.conjoin(torch.cat([x, p], dim=1)) + return self.reduce(x) + + # Similar to ExpansionBlock but does not upsample. class ConjoinBlock(nn.Module): def __init__(self, filters_in, filters_out=None, block=ConvGnSilu, norm=True): diff --git a/codes/models/networks.py b/codes/models/networks.py index 7c7a7c8e..ff536519 100644 --- a/codes/models/networks.py +++ b/codes/models/networks.py @@ -115,6 +115,8 @@ def define_G(opt, net_key='network_G'): netG = spsr.SwitchedSpsr(in_nc=3, out_nc=3, nf=opt_net['nf'], upscale=opt_net['scale']) elif which_model == "spsr_switched_lr": netG = spsr.SwitchedSpsrLr(in_nc=3, out_nc=3, nf=opt_net['nf'], upscale=opt_net['scale']) + elif which_model == "spsr_switched_lr2": + netG = spsr.SwitchedSpsrLr2(in_nc=3, out_nc=3, nf=opt_net['nf'], upscale=opt_net['scale']) # image corruption elif which_model == 'HighToLowResNet':