Move ExpansionBlock to arch_util

Also makes all processing blocks have a conformant signature.

Alters ExpansionBlock to perform a processing conv on the passthrough
before the conjoin operation - this will break backwards compatibilty with SRG2.

codes/models/archs/NestedSwitchGenerator.py

@@ -83,9 +83,9 @@ class Constrictor(nn.Module):
         assert(filters > output_filters)
         gap = filters - output_filters
         gap_div_4 = int(gap / 4)
-        self.cbl1 = ConvBnRelu(filters, filters - (gap_div_4 * 2), kernel_size=1, bn=True, bias=True)
+        self.cbl1 = ConvBnRelu(filters, filters - (gap_div_4 * 2), kernel_size=1, norm=True, bias=True)
-        self.cbl2 = ConvBnRelu(filters - (gap_div_4 * 2), filters - (gap_div_4 * 3), kernel_size=1, bn=True, bias=False)
+        self.cbl2 = ConvBnRelu(filters - (gap_div_4 * 2), filters - (gap_div_4 * 3), kernel_size=1, norm=True, bias=False)
-        self.cbl3 = ConvBnRelu(filters - (gap_div_4 * 3), output_filters, kernel_size=1, relu=False, bn=False, bias=False)
+        self.cbl3 = ConvBnRelu(filters - (gap_div_4 * 3), output_filters, kernel_size=1, activation=False, norm=False, bias=False)
 
     def forward(self, x):
         x = self.cbl1(x)
@@ -134,10 +134,10 @@ class NestedSwitchComputer(nn.Module):
             filters.append(current_filters)
             reduce = True
 
-        self.multiplexer_init_conv = ConvBnLelu(transform_filters, switch_base_filters, kernel_size=7, lelu=False, bn=False)
+        self.multiplexer_init_conv = ConvBnLelu(transform_filters, switch_base_filters, kernel_size=7, activation=False, norm=False)
         self.processing_trunk = nn.ModuleList(processing_trunk)
         self.switch = RecursiveSwitchedTransform(transform_filters, filters, nesting_depth-1, transforms_at_leaf, trans_kernel_size, trans_num_layers-1, trans_scale_init, initial_temp=initial_temp, add_scalable_noise_to_transforms=add_scalable_noise_to_transforms)
-        self.anneal = ConvBnLelu(transform_filters, transform_filters, kernel_size=1, bn=False)
+        self.anneal = ConvBnLelu(transform_filters, transform_filters, kernel_size=1, norm=False)
 
     def forward(self, x):
         feed_forward = x
@@ -161,9 +161,9 @@ class NestedSwitchedGenerator(nn.Module):
                  trans_layers, transformation_filters, 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(NestedSwitchedGenerator, self).__init__()
-        self.initial_conv = ConvBnLelu(3, transformation_filters, kernel_size=7, lelu=False, bn=False)
+        self.initial_conv = ConvBnLelu(3, transformation_filters, kernel_size=7, activation=False, norm=False)
-        self.proc_conv = ConvBnLelu(transformation_filters, transformation_filters, bn=False)
+        self.proc_conv = ConvBnLelu(transformation_filters, transformation_filters, norm=False)
-        self.final_conv = ConvBnLelu(transformation_filters, 3, kernel_size=1, lelu=False, bn=False)
+        self.final_conv = ConvBnLelu(transformation_filters, 3, kernel_size=1, activation=False, norm=False)
 
         switches = []
         for sw_reduce, sw_proc, trans_count, kernel, layers in zip(switch_reductions, switch_processing_layers, trans_counts, trans_kernel_sizes, trans_layers):

codes/models/archs/SRG1_arch.py

@@ -12,9 +12,9 @@ class MultiConvBlock(nn.Module):
         assert depth >= 2
         super(MultiConvBlock, self).__init__()
         self.noise_scale = nn.Parameter(torch.full((1,), fill_value=.01))
-        self.bnconvs = nn.ModuleList([ConvBnLelu(filters_in, filters_mid, kernel_size, bn=bn, bias=False)] +
+        self.bnconvs = nn.ModuleList([ConvBnLelu(filters_in, filters_mid, kernel_size, norm=bn, bias=False)] +
-                                     [ConvBnLelu(filters_mid, filters_mid, kernel_size, bn=bn, bias=False) for i in range(depth-2)] +
+                                     [ConvBnLelu(filters_mid, filters_mid, kernel_size, norm=bn, bias=False) for i in range(depth - 2)] +
-                                     [ConvBnLelu(filters_mid, filters_out, kernel_size, lelu=False, bn=False, bias=False)])
+                                     [ConvBnLelu(filters_mid, filters_out, kernel_size, activation=False, norm=False, bias=False)])
         self.scale = nn.Parameter(torch.full((1,), fill_value=scale_init))
         self.bias = nn.Parameter(torch.zeros(1))
 
@@ -32,8 +32,8 @@ class MultiConvBlock(nn.Module):
 class HalvingProcessingBlock(nn.Module):
     def __init__(self, filters):
         super(HalvingProcessingBlock, self).__init__()
-        self.bnconv1 = ConvBnSilu(filters, filters * 2, stride=2, bn=False, bias=False)
+        self.bnconv1 = ConvBnSilu(filters, filters * 2, stride=2, norm=False, bias=False)
-        self.bnconv2 = ConvBnSilu(filters * 2, filters * 2, bn=True, bias=False)
+        self.bnconv2 = ConvBnSilu(filters * 2, filters * 2, norm=True, bias=False)
     def forward(self, x):
         x = self.bnconv1(x)
         return self.bnconv2(x)
@@ -45,7 +45,7 @@ def create_sequential_growing_processing_block(filters_init, filter_growth, num_
     convs = []
     current_filters = filters_init
     for i in range(num_convs):
-        convs.append(ConvBnSilu(current_filters, current_filters + filter_growth, bn=True, bias=False))
+        convs.append(ConvBnSilu(current_filters, current_filters + filter_growth, norm=True, bias=False))
         current_filters += filter_growth
     return nn.Sequential(*convs), current_filters
 
@@ -60,7 +60,7 @@ class SwitchComputer(nn.Module):
         self.reduction_blocks = nn.ModuleList([HalvingProcessingBlock(filters * 2 ** i) for i in range(reduction_blocks)])
         final_filters = filters * 2 ** reduction_blocks
         self.processing_blocks, final_filters = create_sequential_growing_processing_block(final_filters, growth, processing_blocks)
-        self.post_interpolate_decimate = ConvBnSilu(final_filters, filters, kernel_size=1, silu=False, bn=False)
+        self.post_interpolate_decimate = ConvBnSilu(final_filters, filters, kernel_size=1, activation=False, norm=False)
         self.interpolate_process = ConvBnSilu(filters, filters)
         self.interpolate_process2 = ConvBnSilu(filters, filters)
         tc = transform_count

codes/models/archs/SwitchedResidualGenerator_arch.py

@@ -4,7 +4,7 @@ from switched_conv import BareConvSwitch, compute_attention_specificity
 import torch.nn.functional as F
 import functools
 from collections import OrderedDict
-from models.archs.arch_util import ConvBnLelu, ConvGnSilu
+from models.archs.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock
 from models.archs.RRDBNet_arch import ResidualDenseBlock_5C
 from models.archs.spinenet_arch import SpineNet
 from switched_conv_util import save_attention_to_image
@@ -15,9 +15,9 @@ class MultiConvBlock(nn.Module):
         assert depth >= 2
         super(MultiConvBlock, self).__init__()
         self.noise_scale = nn.Parameter(torch.full((1,), fill_value=.01))
-        self.bnconvs = nn.ModuleList([ConvBnLelu(filters_in, filters_mid, kernel_size, bn=bn, bias=False, weight_init_factor=weight_init_factor)] +
+        self.bnconvs = nn.ModuleList([ConvBnLelu(filters_in, filters_mid, kernel_size, norm=bn, bias=False, weight_init_factor=weight_init_factor)] +
-                                     [ConvBnLelu(filters_mid, filters_mid, kernel_size, bn=bn, bias=False, weight_init_factor=weight_init_factor) for i in range(depth-2)] +
+                                     [ConvBnLelu(filters_mid, filters_mid, kernel_size, norm=bn, bias=False, weight_init_factor=weight_init_factor) for i in range(depth - 2)] +
-                                     [ConvBnLelu(filters_mid, filters_out, kernel_size, lelu=False, bn=False, bias=False, weight_init_factor=weight_init_factor)])
+                                     [ConvBnLelu(filters_mid, filters_out, kernel_size, activation=False, norm=False, bias=False, weight_init_factor=weight_init_factor)])
         self.scale = nn.Parameter(torch.full((1,), fill_value=scale_init))
         self.bias = nn.Parameter(torch.zeros(1))
 
@@ -35,28 +35,14 @@ class MultiConvBlock(nn.Module):
 class HalvingProcessingBlock(nn.Module):
     def __init__(self, filters):
         super(HalvingProcessingBlock, self).__init__()
-        self.bnconv1 = ConvGnSilu(filters, filters * 2, stride=2, gn=False, bias=False)
+        self.bnconv1 = ConvGnSilu(filters, filters * 2, stride=2, norm=False, bias=False)
-        self.bnconv2 = ConvGnSilu(filters * 2, filters * 2, gn=True, bias=False)
+        self.bnconv2 = ConvGnSilu(filters * 2, filters * 2, norm=True, bias=False)
 
     def forward(self, x):
         x = self.bnconv1(x)
         return self.bnconv2(x)
 
 
-class ExpansionBlock(nn.Module):
-    def __init__(self, filters):
-        super(ExpansionBlock, self).__init__()
-        self.decimate = ConvGnSilu(filters, filters // 2, kernel_size=1, bias=False, silu=False, gn=False)
-        self.conjoin = ConvGnSilu(filters, filters // 2, kernel_size=3, bias=True, silu=False, gn=True)
-        self.process = ConvGnSilu(filters // 2, filters // 2, kernel_size=3, bias=False, silu=True, gn=True)
-
-    def forward(self, input, passthrough):
-        x = F.interpolate(input, scale_factor=2, mode="nearest")
-        x = self.decimate(x)
-        x = self.conjoin(torch.cat([x, passthrough], dim=1))
-        return self.process(x)
-
-
 # This is a classic u-net architecture with the goal of assigning each individual pixel an individual transform
 # switching set.
 class ConvBasisMultiplexer(nn.Module):
@@ -70,10 +56,10 @@ class ConvBasisMultiplexer(nn.Module):
 
         gap = base_filters - multiplexer_channels
         cbl1_out = ((base_filters - (gap // 2)) // 4) * 4   # Must be multiples of 4 to use with group norm.
-        self.cbl1 = ConvGnSilu(base_filters, cbl1_out, gn=use_gn, bias=False, num_groups=4)
+        self.cbl1 = ConvGnSilu(base_filters, cbl1_out, norm=use_gn, bias=False, num_groups=4)
         cbl2_out = ((base_filters - (3 * gap // 4)) // 4) * 4
-        self.cbl2 = ConvGnSilu(cbl1_out, cbl2_out, gn=use_gn, bias=False, num_groups=4)
+        self.cbl2 = ConvGnSilu(cbl1_out, cbl2_out, norm=use_gn, bias=False, num_groups=4)
-        self.cbl3 = ConvGnSilu(cbl2_out, multiplexer_channels, bias=True, gn=False)
+        self.cbl3 = ConvGnSilu(cbl2_out, multiplexer_channels, bias=True, norm=False)
 
     def forward(self, x):
         x = self.filter_conv(x)
@@ -109,11 +95,11 @@ class BackboneMultiplexer(nn.Module):
         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, gn=False, silu=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, gn=False, silu=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, gn=False, silu=False)
+        self.final = ConvGnSilu(64, transform_count, bias=False, norm=False, activation=False)
 
     def forward(self, x):
         spine = self.backbone.get_forward_result()
@@ -139,7 +125,7 @@ class ConfigurableSwitchComputer(nn.Module):
         # And the switch itself, including learned scalars
         self.switch = BareConvSwitch(initial_temperature=init_temp)
         self.switch_scale = nn.Parameter(torch.full((1,), float(1)))
-        self.post_switch_conv = ConvBnLelu(base_filters, base_filters, bn=False, bias=True)
+        self.post_switch_conv = ConvBnLelu(base_filters, base_filters, norm=False, bias=True)
         # The post_switch_conv gets a low scale initially. The network can decide to magnify it (or not)
         # depending on its needs.
         self.psc_scale = nn.Parameter(torch.full((1,), float(.1)))
@@ -174,11 +160,11 @@ class ConfigurableSwitchedResidualGenerator2(nn.Module):
                  add_scalable_noise_to_transforms=False):
         super(ConfigurableSwitchedResidualGenerator2, self).__init__()
         switches = []
-        self.initial_conv = ConvBnLelu(3, transformation_filters, bn=False, lelu=False, bias=True)
+        self.initial_conv = ConvBnLelu(3, transformation_filters, norm=False, activation=False, bias=True)
-        self.upconv1 = ConvBnLelu(transformation_filters, transformation_filters, bn=False, bias=True)
+        self.upconv1 = ConvBnLelu(transformation_filters, transformation_filters, norm=False, bias=True)
-        self.upconv2 = ConvBnLelu(transformation_filters, transformation_filters, bn=False, bias=True)
+        self.upconv2 = ConvBnLelu(transformation_filters, transformation_filters, norm=False, bias=True)
-        self.hr_conv = ConvBnLelu(transformation_filters, transformation_filters, bn=False, bias=True)
+        self.hr_conv = ConvBnLelu(transformation_filters, transformation_filters, norm=False, bias=True)
-        self.final_conv = ConvBnLelu(transformation_filters, 3, bn=False, lelu=False, bias=True)
+        self.final_conv = ConvBnLelu(transformation_filters, 3, norm=False, activation=False, bias=True)
         for _ in range(switch_depth):
             multiplx_fn = functools.partial(ConvBasisMultiplexer, transformation_filters, switch_filters, switch_reductions, switch_processing_layers, trans_counts)
             pretransform_fn = functools.partial(ConvBnLelu, transformation_filters, transformation_filters, bn=False, bias=False, weight_init_factor=.1)
@@ -258,19 +244,19 @@ class ConfigurableSwitchedResidualGenerator3(nn.Module):
                  heightened_temp_min=1,
                  heightened_final_step=50000, upsample_factor=4):
         super(ConfigurableSwitchedResidualGenerator3, self).__init__()
-        self.initial_conv = ConvBnLelu(3, base_filters, bn=False, lelu=False, bias=True)
+        self.initial_conv = ConvBnLelu(3, base_filters, norm=False, activation=False, bias=True)
-        self.sw_conv = ConvBnLelu(base_filters, base_filters, lelu=False, bias=True)
+        self.sw_conv = ConvBnLelu(base_filters, base_filters, activation=False, bias=True)
-        self.upconv1 = ConvBnLelu(base_filters, base_filters, bn=False, bias=True)
+        self.upconv1 = ConvBnLelu(base_filters, base_filters, norm=False, bias=True)
-        self.upconv2 = ConvBnLelu(base_filters, base_filters, bn=False, bias=True)
+        self.upconv2 = ConvBnLelu(base_filters, base_filters, norm=False, bias=True)
-        self.hr_conv = ConvBnLelu(base_filters, base_filters, bn=False, bias=True)
+        self.hr_conv = ConvBnLelu(base_filters, base_filters, norm=False, bias=True)
-        self.final_conv = ConvBnLelu(base_filters, 3, bn=False, lelu=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, bn=False, lelu=False, bias=False))
+        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)

codes/models/archs/arch_util.py

@@ -184,16 +184,16 @@ class SiLU(nn.Module):
 ''' Convenience class with Conv->BN->ReLU. Includes weight initialization and auto-padding for standard
     kernel sizes. '''
 class ConvBnRelu(nn.Module):
-    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, relu=True, bn=True, bias=True):
+    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, activation=True, norm=True, bias=True):
         super(ConvBnRelu, self).__init__()
         padding_map = {1: 0, 3: 1, 5: 2, 7: 3}
         assert kernel_size in padding_map.keys()
         self.conv = nn.Conv2d(filters_in, filters_out, kernel_size, stride, padding_map[kernel_size], bias=bias)
-        if bn:
+        if norm:
             self.bn = nn.BatchNorm2d(filters_out)
         else:
             self.bn = None
-        if relu:
+        if activation:
             self.relu = nn.ReLU()
         else:
             self.relu = None
@@ -219,16 +219,16 @@ class ConvBnRelu(nn.Module):
 ''' Convenience class with Conv->BN->SiLU. Includes weight initialization and auto-padding for standard
     kernel sizes. '''
 class ConvBnSilu(nn.Module):
-    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, silu=True, bn=True, bias=True, weight_init_factor=1):
+    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, activation=True, norm=True, bias=True, weight_init_factor=1):
         super(ConvBnSilu, self).__init__()
         padding_map = {1: 0, 3: 1, 5: 2, 7: 3}
         assert kernel_size in padding_map.keys()
         self.conv = nn.Conv2d(filters_in, filters_out, kernel_size, stride, padding_map[kernel_size], bias=bias)
-        if bn:
+        if norm:
             self.bn = nn.BatchNorm2d(filters_out)
         else:
             self.bn = None
-        if silu:
+        if activation:
             self.silu = SiLU()
         else:
             self.silu = None
@@ -257,16 +257,16 @@ class ConvBnSilu(nn.Module):
 ''' Convenience class with Conv->BN->LeakyReLU. Includes weight initialization and auto-padding for standard
     kernel sizes. '''
 class ConvBnLelu(nn.Module):
-    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, lelu=True, bn=True, bias=True, weight_init_factor=1):
+    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, activation=True, norm=True, bias=True, weight_init_factor=1):
         super(ConvBnLelu, self).__init__()
         padding_map = {1: 0, 3: 1, 5: 2, 7: 3}
         assert kernel_size in padding_map.keys()
         self.conv = nn.Conv2d(filters_in, filters_out, kernel_size, stride, padding_map[kernel_size], bias=bias)
-        if bn:
+        if norm:
             self.bn = nn.BatchNorm2d(filters_out)
         else:
             self.bn = None
-        if lelu:
+        if activation:
             self.lelu = nn.LeakyReLU(negative_slope=.1)
         else:
             self.lelu = None
@@ -296,16 +296,16 @@ class ConvBnLelu(nn.Module):
 ''' Convenience class with Conv->GroupNorm->LeakyReLU. Includes weight initialization and auto-padding for standard
     kernel sizes. '''
 class ConvGnLelu(nn.Module):
-    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, lelu=True, gn=True, bias=True, num_groups=8):
+    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, activation=True, norm=True, bias=True, num_groups=8):
         super(ConvGnLelu, self).__init__()
         padding_map = {1: 0, 3: 1, 5: 2, 7: 3}
         assert kernel_size in padding_map.keys()
         self.conv = nn.Conv2d(filters_in, filters_out, kernel_size, stride, padding_map[kernel_size], bias=bias)
-        if gn:
+        if norm:
             self.gn = nn.GroupNorm(num_groups, filters_out)
         else:
             self.gn = None
-        if lelu:
+        if activation:
             self.lelu = nn.LeakyReLU(negative_slope=.1)
         else:
             self.lelu = None
@@ -331,16 +331,16 @@ class ConvGnLelu(nn.Module):
 ''' Convenience class with Conv->BN->SiLU. Includes weight initialization and auto-padding for standard
     kernel sizes. '''
 class ConvGnSilu(nn.Module):
-    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, silu=True, gn=True, bias=True, num_groups=8, weight_init_factor=1):
+    def __init__(self, filters_in, filters_out, kernel_size=3, stride=1, activation=True, norm=True, bias=True, num_groups=8, weight_init_factor=1):
         super(ConvGnSilu, self).__init__()
         padding_map = {1: 0, 3: 1, 5: 2, 7: 3}
         assert kernel_size in padding_map.keys()
         self.conv = nn.Conv2d(filters_in, filters_out, kernel_size, stride, padding_map[kernel_size], bias=bias)
-        if gn:
+        if norm:
             self.gn = nn.GroupNorm(num_groups, filters_out)
         else:
             self.gn = None
-        if silu:
+        if activation:
             self.silu = SiLU()
         else:
             self.silu = None
@@ -363,4 +363,24 @@ class ConvGnSilu(nn.Module):
         if self.silu:
             return self.silu(x)
         else:
-            return x
+            return x
+
+# Block that upsamples 2x and reduces incoming filters by 2x. It preserves structure by taking a passthrough feed
+# along with the feature representation.
+class ExpansionBlock(nn.Module):
+    def __init__(self, filters, block=ConvGnSilu):
+        super(ExpansionBlock, self).__init__()
+        self.decimate = block(filters, filters // 2, kernel_size=1, bias=False, activation=False, norm=True)
+        self.process_passthrough = block(filters // 2, filters // 2, kernel_size=3, bias=True, activation=False, norm=True)
+        self.conjoin = block(filters, filters // 2, kernel_size=3, bias=False, activation=True, norm=False)
+        self.process = block(filters // 2, filters // 2, 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.process(x)

codes/models/archs/discriminator_vgg_arch.py

@@ -108,20 +108,20 @@ class Discriminator_VGG_PixLoss(nn.Module):
         self.bn4_1 = nn.GroupNorm(8, nf * 8, affine=True)
 
         self.reduce_1 = ConvGnLelu(nf * 8, nf * 4, bias=False)
-        self.pix_loss_collapse = ConvGnLelu(nf * 4, 1, bias=False, gn=False, lelu=False)
+        self.pix_loss_collapse = ConvGnLelu(nf * 4, 1, bias=False, norm=False, activation=False)
 
         # Pyramid network: upsample with residuals and produce losses at multiple resolutions.
-        self.up3_decimate = ConvGnLelu(nf * 8, nf * 8, kernel_size=3, bias=True, lelu=False)
+        self.up3_decimate = ConvGnLelu(nf * 8, nf * 8, kernel_size=3, bias=True, activation=False)
         self.up3_converge = ConvGnLelu(nf * 16, nf * 8, kernel_size=3, bias=False)
         self.up3_proc = ConvGnLelu(nf * 8, nf * 8, bias=False)
         self.up3_reduce = ConvGnLelu(nf * 8, nf * 4, bias=False)
-        self.up3_pix = ConvGnLelu(nf * 4, 1, bias=False, gn=False, lelu=False)
+        self.up3_pix = ConvGnLelu(nf * 4, 1, bias=False, norm=False, activation=False)
 
-        self.up2_decimate = ConvGnLelu(nf * 8, nf * 4, kernel_size=1, bias=True, lelu=False)
+        self.up2_decimate = ConvGnLelu(nf * 8, nf * 4, kernel_size=1, bias=True, activation=False)
         self.up2_converge = ConvGnLelu(nf * 8, nf * 4, kernel_size=3, bias=False)
         self.up2_proc = ConvGnLelu(nf * 4, nf * 4, bias=False)
         self.up2_reduce = ConvGnLelu(nf * 4, nf * 2, bias=False)
-        self.up2_pix = ConvGnLelu(nf * 2, 1, bias=False, gn=False, lelu=False)
+        self.up2_pix = ConvGnLelu(nf * 2, 1, bias=False, norm=False, activation=False)
 
         # activation function
         self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)