diff --git a/codes/models/archs/SPSR_arch.py b/codes/models/archs/SPSR_arch.py index 414b43ff..a674ea52 100644 --- a/codes/models/archs/SPSR_arch.py +++ b/codes/models/archs/SPSR_arch.py @@ -473,127 +473,3 @@ class SwitchedSpsrWithRef(nn.Module): val["switch_%i_specificity" % (i,)] = means[i] val["switch_%i_histogram" % (i,)] = hists[i] return val - - -class SwitchedSpsrWithRef4x(nn.Module): - def __init__(self, in_nc, out_nc, nf, xforms=8, upscale=4, init_temperature=10): - super(SwitchedSpsrWithRef4x, self).__init__() - n_upscale = int(math.log(upscale, 2)) - - # switch options - transformation_filters = nf - switch_filters = nf - self.transformation_counts = xforms - self.reference_processor = ReferenceImageBranch(transformation_filters) - multiplx_fn = functools.partial(ReferencingConvMultiplexer, transformation_filters, switch_filters, self.transformation_counts) - pretransform_fn = functools.partial(AdaInConvBlock, 512, transformation_filters, transformation_filters) - 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=init_temperature, - 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=init_temperature, - add_scalable_noise_to_transforms=True) - self.feature_lr_conv = ConvGnLelu(nf, nf, kernel_size=3, norm=True, activation=False) - self.stage1_up_fea = UpconvBlock(nf, nf, block=ConvGnLelu, norm=False, activation=False, 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) - mplex_grad = functools.partial(ReferencingConvMultiplexer, nf * 2, nf * 2, self.transformation_counts // 2) - 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=init_temperature, - add_scalable_noise_to_transforms=True) - self.stage1_up_grad = UpconvBlock(nf, nf, block=ConvGnLelu, norm=False, activation=False, bias=False) - - # 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. - 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(AdaInConvBlock, 512, transformation_filters * 2, transformation_filters * 2) - 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=init_temperature, - add_scalable_noise_to_transforms=True) - self.stage2_up_fea = UpconvBlock(nf, nf, block=ConvGnLelu, norm=False, activation=False, bias=False) - self.stage2_up_grad = UpconvBlock(nf, nf, block=ConvGnLelu, norm=False, activation=False, bias=False) - 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 = init_temperature - self.final_temperature_step = 10000 - - def forward(self, x, ref, center_coord): - x_grad = self.get_g_nopadding(x) - ref = self.reference_processor(ref, center_coord) - x = self.model_fea_conv(x) - - x1, a1 = self.sw1((x, ref), True) - x2, a2 = self.sw2((x1, ref), True) - x_fea = self.feature_lr_conv(x2) - x_fea = self.stage1_up_fea(x_fea) - 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, ref), att_in=(torch.cat([x1, x_b_fea], dim=1), ref), output_attention_weights=True) - x_grad = self.grad_lr_conv(x_grad) - x_grad = self.stage1_up_grad(x_grad) - x_grad = self.grad_hr_conv(x_grad) - x_out_branch = self.stage2_up_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, ref), att_in=(x_fea, ref), identity=x_fea, output_attention_weights=True) - x_out = self.final_lr_conv(x__branch_pretrain_cat) - x_out = self.stage2_up_fea(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 4f0bac56..988e88cb 100644 --- a/codes/models/archs/SwitchedResidualGenerator_arch.py +++ b/codes/models/archs/SwitchedResidualGenerator_arch.py @@ -4,9 +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, ExpansionBlock2, ConjoinBlock, ConvGnLelu -from models.archs.RRDBNet_arch import ResidualDenseBlock_5C, RRDB -from models.archs.spinenet_arch import SpineNet +from models.archs.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock, ExpansionBlock2, ConvGnLelu from switched_conv_util import save_attention_to_image_rgb import os @@ -81,17 +79,6 @@ class ConvBasisMultiplexer(nn.Module): return x -class CachedBackboneWrapper: - def __init__(self, backbone: nn.Module): - self.backbone = backbone - - def __call__(self, *args): - self.cache = self.backbone(*args) - return self.cache - - def get_forward_result(self): - return self.cache - # torch.gather() which operates across 2d images. def gather_2d(input, index): b, c, h, w = input.shape @@ -187,26 +174,6 @@ class ReferencingConvMultiplexer(nn.Module): return x -class BackboneMultiplexer(nn.Module): - def __init__(self, backbone: CachedBackboneWrapper, transform_count): - super(BackboneMultiplexer, self).__init__() - self.backbone = backbone - self.proc = nn.Sequential(ConvGnSilu(256, 256, kernel_size=3, bias=True), - ConvGnSilu(256, 256, kernel_size=3, bias=False)) - self.up1 = nn.Sequential(ConvGnSilu(256, 128, kernel_size=3, bias=False, norm=False, activation=False), - ConvGnSilu(128, 128, kernel_size=3, bias=False)) - self.up2 = nn.Sequential(ConvGnSilu(128, 64, kernel_size=3, bias=False, norm=False, activation=False), - ConvGnSilu(64, 64, kernel_size=3, bias=False)) - self.final = ConvGnSilu(64, transform_count, bias=False, norm=False, activation=False) - - def forward(self, x): - spine = self.backbone.get_forward_result() - feat = self.proc(spine[0]) - feat = self.up1(F.interpolate(feat, scale_factor=2, mode="nearest")) - feat = self.up2(F.interpolate(feat, scale_factor=2, mode="nearest")) - return self.final(feat) - - class ConfigurableSwitchComputer(nn.Module): def __init__(self, base_filters, multiplexer_net, pre_transform_block, transform_block, transform_count, attention_norm, init_temp=20, add_scalable_noise_to_transforms=False): @@ -364,204 +331,3 @@ class ConfigurableSwitchedResidualGenerator2(nn.Module): val["switch_%i_histogram" % (i,)] = hists[i] return val - -# Equivalent to SRG2 - Uses RDB blocks in between two switches. -class ConfigurableSwitchedResidualGenerator4(nn.Module): - def __init__(self, switch_filters, switch_reductions, switch_processing_layers, trans_counts, trans_kernel_sizes, - trans_layers, transformation_filters, attention_norm, initial_temp=20, final_temperature_step=50000, heightened_temp_min=1, - heightened_final_step=50000, upsample_factor=1, - add_scalable_noise_to_transforms=False): - super(ConfigurableSwitchedResidualGenerator4, self).__init__() - self.initial_conv = ConvBnLelu(3, transformation_filters, norm=False, activation=False, bias=True) - self.upconv1 = ConvBnLelu(transformation_filters, transformation_filters, norm=False, bias=True) - self.upconv2 = ConvBnLelu(transformation_filters, transformation_filters, norm=False, bias=True) - self.hr_conv = ConvBnLelu(transformation_filters, transformation_filters, norm=False, bias=True) - - multiplx_fn = functools.partial(ConvBasisMultiplexer, transformation_filters, switch_filters, switch_reductions, - switch_processing_layers, trans_counts) - half_multiplx_fn = functools.partial(ConvBasisMultiplexer, transformation_filters, switch_filters, switch_reductions, - switch_processing_layers, trans_counts // 2) - transform_fn = functools.partial(MultiConvBlock, transformation_filters, int(transformation_filters * 1.5), - transformation_filters, kernel_size=trans_kernel_sizes, depth=trans_layers, - weight_init_factor=.1) - self.rdb1 = RRDB(transformation_filters) - self.sw1 = ConfigurableSwitchComputer(transformation_filters, multiplx_fn, - pre_transform_block=None, transform_block=transform_fn, - attention_norm=attention_norm, - transform_count=trans_counts, init_temp=initial_temp, - add_scalable_noise_to_transforms=add_scalable_noise_to_transforms) - self.rdb2 = RRDB(transformation_filters) - self.sw2 = ConfigurableSwitchComputer(transformation_filters, half_multiplx_fn, - pre_transform_block=None, transform_block=transform_fn, - attention_norm=attention_norm, - transform_count=trans_counts // 2, init_temp=initial_temp, - add_scalable_noise_to_transforms=add_scalable_noise_to_transforms) - self.rdb3 = RRDB(transformation_filters) - self.sw3 = ConfigurableSwitchComputer(transformation_filters, multiplx_fn, - pre_transform_block=None, transform_block=transform_fn, - attention_norm=attention_norm, - transform_count=trans_counts, init_temp=initial_temp, - add_scalable_noise_to_transforms=add_scalable_noise_to_transforms) - self.rdb4 = RRDB(transformation_filters) - self.switches = [self.sw1, self.sw2, self.sw3] - - self.final_conv = ConvBnLelu(transformation_filters, 3, norm=False, activation=False, bias=True) - self.transformation_counts = trans_counts - self.init_temperature = initial_temp - self.final_temperature_step = final_temperature_step - self.heightened_temp_min = heightened_temp_min - self.heightened_final_step = heightened_final_step - self.attentions = None - self.upsample_factor = upsample_factor - assert self.upsample_factor == 2 or self.upsample_factor == 4 - - def forward(self, x): - # This is a common bug when evaluating SRG2 generators. It needs to be configured properly in eval mode. Just fail. - if not self.train: - assert self.switches[0].switch.temperature == 1 - - x = self.initial_conv(x) - - x = self.rdb1(x) - x, a1 = self.sw1(x, True) - x = self.rdb2(x) - x, a2 = self.sw2(x, True) - x = self.rdb3(x) - x, a3 = self.sw3(x, True) - x = self.rdb4(x) - self.attentions = [a1, a2, a3] - - x = self.upconv1(F.interpolate(x, scale_factor=2, mode="nearest")) - if self.upsample_factor > 2: - x = F.interpolate(x, scale_factor=2, mode="nearest") - x = self.upconv2(x) - x = self.final_conv(self.hr_conv(x)) - return x, x - - 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) - if temp == 1 and self.heightened_final_step and step > self.final_temperature_step and \ - self.heightened_final_step != 1: - # Once the temperature passes (1) it enters an inverted curve to match the linear curve from above. - # without this, the attention specificity "spikes" incredibly fast in the last few iterations. - h_steps_total = self.heightened_final_step - self.final_temperature_step - h_steps_current = min(step - self.final_temperature_step, h_steps_total) - # The "gap" will represent the steps that need to be traveled as a linear function. - h_gap = 1 / self.heightened_temp_min - temp = h_gap * h_steps_current / h_steps_total - # Invert temperature to represent reality on this side of the curve - temp = 1 / temp - self.set_temperature(temp) - if step % 50 == 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 - -class Interpolate(nn.Module): - def __init__(self, factor): - super(Interpolate, self).__init__() - self.factor = factor - - def forward(self, x): - return F.interpolate(x, scale_factor=self.factor) - - -class ConfigurableSwitchedResidualGenerator3(nn.Module): - def __init__(self, base_filters, trans_count, initial_temp=20, final_temperature_step=50000, - heightened_temp_min=1, - heightened_final_step=50000, upsample_factor=4): - super(ConfigurableSwitchedResidualGenerator3, self).__init__() - self.initial_conv = ConvBnLelu(3, base_filters, norm=False, activation=False, bias=True) - self.sw_conv = ConvBnLelu(base_filters, base_filters, activation=False, bias=True) - self.upconv1 = ConvBnLelu(base_filters, base_filters, norm=False, bias=True) - self.upconv2 = ConvBnLelu(base_filters, base_filters, norm=False, bias=True) - self.hr_conv = ConvBnLelu(base_filters, base_filters, norm=False, bias=True) - self.final_conv = ConvBnLelu(base_filters, 3, norm=False, activation=False, bias=True) - - self.backbone = SpineNet('49', in_channels=3, use_input_norm=True) - for p in self.backbone.parameters(recurse=True): - p.requires_grad = False - self.backbone_wrapper = CachedBackboneWrapper(self.backbone) - multiplx_fn = functools.partial(BackboneMultiplexer, self.backbone_wrapper) - pretransform_fn = functools.partial(nn.Sequential, ConvBnLelu(base_filters, base_filters, kernel_size=3, norm=False, activation=False, bias=False)) - transform_fn = functools.partial(MultiConvBlock, base_filters, int(base_filters * 1.5), base_filters, kernel_size=3, depth=4) - self.switch = ConfigurableSwitchComputer(base_filters, multiplx_fn, pretransform_fn, transform_fn, trans_count, init_temp=initial_temp, - add_scalable_noise_to_transforms=True, init_scalar=.1) - - self.transformation_counts = trans_count - self.init_temperature = initial_temp - self.final_temperature_step = final_temperature_step - self.heightened_temp_min = heightened_temp_min - self.heightened_final_step = heightened_final_step - self.attentions = None - self.upsample_factor = upsample_factor - self.backbone_forward = None - - def get_forward_results(self): - return self.backbone_forward - - def forward(self, x): - self.backbone_forward = self.backbone_wrapper(F.interpolate(x, scale_factor=2, mode="nearest")) - - x = self.initial_conv(x) - - self.attentions = [] - x, att = self.switch(x, output_attention_weights=True) - self.attentions.append(att) - - x = self.upconv1(F.interpolate(x, scale_factor=2, mode="nearest")) - if self.upsample_factor > 2: - x = F.interpolate(x, scale_factor=2, mode="nearest") - x = self.upconv2(x) - return self.final_conv(self.hr_conv(x)), - - def set_temperature(self, temp): - self.switch.set_temperature(temp) - - 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) - if temp == 1 and self.heightened_final_step and step > self.final_temperature_step and \ - self.heightened_final_step != 1: - # Once the temperature passes (1) it enters an inverted curve to match the linear curve from above. - # without this, the attention specificity "spikes" incredibly fast in the last few iterations. - h_steps_total = self.heightened_final_step - self.final_temperature_step - h_steps_current = min(step - self.final_temperature_step, h_steps_total) - # The "gap" will represent the steps that need to be traveled as a linear function. - h_gap = 1 / self.heightened_temp_min - temp = h_gap * h_steps_current / h_steps_total - # Invert temperature to represent reality on this side of the curve - temp = 1 / temp - self.set_temperature(temp) - if step % 50 == 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.switch.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/networks.py b/codes/models/networks.py index d9740ab4..f25ce61b 100644 --- a/codes/models/networks.py +++ b/codes/models/networks.py @@ -50,21 +50,6 @@ def define_G(opt, net_key='network_G', scale=None): initial_temp=opt_net['temperature'], final_temperature_step=opt_net['temperature_final_step'], heightened_temp_min=opt_net['heightened_temp_min'], heightened_final_step=opt_net['heightened_final_step'], upsample_factor=scale, add_scalable_noise_to_transforms=opt_net['add_noise']) - elif which_model == "ConfigurableSwitchedResidualGenerator4": - netG = SwitchedGen_arch.ConfigurableSwitchedResidualGenerator4(switch_filters=opt_net['switch_filters'], - switch_reductions=opt_net['switch_reductions'], - switch_processing_layers=opt_net['switch_processing_layers'], trans_counts=opt_net['trans_counts'], - trans_kernel_sizes=opt_net['trans_kernel_sizes'], trans_layers=opt_net['trans_layers'], - transformation_filters=opt_net['transformation_filters'], attention_norm=opt_net['attention_norm'], - initial_temp=opt_net['temperature'], final_temperature_step=opt_net['temperature_final_step'], - heightened_temp_min=opt_net['heightened_temp_min'], heightened_final_step=opt_net['heightened_final_step'], - upsample_factor=scale, add_scalable_noise_to_transforms=opt_net['add_noise']) - elif which_model == 'spsr_net': - netG = spsr.SPSRNet(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', bl_inc=opt_net['bl_inc']) - 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']) @@ -78,19 +63,8 @@ def define_G(opt, net_key='network_G', scale=None): init_temperature=opt_net['temperature'] if 'temperature' in opt_net.keys() else 10) elif which_model == "spsr_switched_with_ref4x": xforms = opt_net['num_transforms'] if 'num_transforms' in opt_net.keys() else 8 - netG = spsr.SwitchedSpsrWithRef4x(in_nc=3, out_nc=3, nf=opt_net['nf'], xforms=xforms, upscale=opt_net['scale'], + netG = spsr.SwitchedSpsrWithRef4x(in_nc=3, out_nc=3, nf=opt_net['nf'], xforms=xforms, init_temperature=opt_net['temperature'] if 'temperature' in opt_net.keys() else 10) - - # image corruption - elif which_model == 'HighToLowResNet': - netG = HighToLowResNet.HighToLowResNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], - nf=opt_net['nf'], nb=opt_net['nb'], downscale=opt_net['scale']) - elif which_model == 'FlatProcessorNet': - '''netG = FlatProcessorNet_arch.FlatProcessorNet(in_nc=opt_net['in_nc'], out_nc=opt_net['out_nc'], - nf=opt_net['nf'], downscale=opt_net['scale'], reduce_anneal_blocks=opt_net['ra_blocks'], - assembler_blocks=opt_net['assembler_blocks'])''' - netG = FlatProcessorNetNew_arch.fixup_resnet34(num_filters=opt_net['nf'])\ - else: raise NotImplementedError('Generator model [{:s}] not recognized'.format(which_model))