Update archs/network.py

archs/network.py

@@ -3,10 +3,10 @@ import torch.nn as nn
 import torch.nn.functional as F
 import functools
 try:
-    from .arch_util import EBlock, Attention_Light
+    from .arch_util import EBlock
     from .arch_util_freq import EBlock_freq
 except:
-    from arch_util import EBlock, Attention_Light
+    from arch_util import EBlock
     from arch_util_freq import EBlock_freq
 
 
@@ -14,11 +14,13 @@ class Network(nn.Module):
     
     def __init__(self, img_channel=3, 
                  width=16, 
-                 middle_blk_num=1, 
+                 middle_blk_num_enc=1,
+                 middle_blk_num_dec=1, 
                  enc_blk_nums=[], 
                  dec_blk_nums=[],  
                  dilations = [1], 
-                 extra_depth_wise = False):
+                 extra_depth_wise = False,
+                 ksize = 5):
         super(Network, self).__init__()
         
         self.intro = nn.Conv2d(in_channels=img_channel, out_channels=width, kernel_size=3, padding=1, stride=1, groups=1,
@@ -36,7 +38,7 @@ class Network(nn.Module):
         for num in enc_blk_nums:
             self.encoders.append(
                 nn.Sequential(
-                    *[EBlock(chan, dilations = dilations, extra_depth_wise=extra_depth_wise) for _ in range(num)]
+                    *[EBlock_freq(chan, extra_depth_wise=extra_depth_wise) for _ in range(num)]
                 )
             )
             self.downs.append(
@@ -44,9 +46,13 @@ class Network(nn.Module):
             )
             chan = chan * 2
 
-        self.middle_blks = \
+        self.middle_blks_enc = \
             nn.Sequential(
-                *[EBlock(chan, dilations = dilations, extra_depth_wise=extra_depth_wise) for _ in range(middle_blk_num)]
+                *[EBlock_freq(chan, extra_depth_wise=extra_depth_wise) for _ in range(middle_blk_num_enc)]
+            )
+        self.middle_blks_dec = \
+            nn.Sequential(
+                *[EBlock(chan, dilations = dilations, extra_depth_wise=extra_depth_wise) for _ in range(middle_blk_num_dec)]
             )
 
         for num in dec_blk_nums:
@@ -59,21 +65,16 @@ class Network(nn.Module):
             chan = chan // 2
             self.decoders.append(
                 nn.Sequential(
-                    *[EBlock(chan, extra_depth_wise=extra_depth_wise) for _ in range(num)]
+                    *[EBlock(chan,dilations = dilations, extra_depth_wise=extra_depth_wise) for _ in range(num)]
                 )
             )
 
         self.padder_size = 2 ** len(self.encoders)        
         
-        #define the attention layers 
-        
-        # self.recon_trunk_light = nn.Sequential(*[FBlock(c = chan * self.padder_size,
-        #                                         DW_Expand=2, FFN_Expand=2, dilations = dilations, 
-        #                                         extra_depth_wise = False) for i in range(residual_layers)])
-
-        # ResidualBlock_noBN_f = functools.partial(ResidualBlock_noBN, nf = width * self.padder_size)
-        # self.recon_trunk_light = make_layer(ResidualBlock_noBN_f, residual_layers)
-        
+        # self.facs = nn.ModuleList([nn.Identity(), nn.Identity(),
+        #                           nn.Identity(),
+        #                           nn.Identity())
+        # self.kconv_deblur = KernelConv2D(ksize=ksize, act = True)
    
         
     def forward(self, input):
@@ -83,26 +84,43 @@ class Network(nn.Module):
         input = self.check_image_size(input)
         x = self.intro(input)
         
-        encs = []
+        # encs = []
+        facs = []
         # i = 0
         for encoder, down in zip(self.encoders, self.downs):
             x = encoder(x)
+            # x_fac = fac(x)
+            facs.append(x)
             # print(i, x.shape)
-            encs.append(x)
+            # encs.append(x)
             x = down(x)
             # i += 1
 
-        x = self.middle_blks(x)
+        # we apply the encoder transforms
+        x_light = self.middle_blks_enc(x)
+        # calculate the fac at this level
+        # x_fac = self.facs[-1](x)
+        # facs.append(x_fac)
+        # apply the decoder transforms
+        x = self.middle_blks_dec(x_light)
+        # apply the fac transform over this step
+        x = x + x_light
+
         # print('3', x.shape)
         # apply the mask
         # x = x * mask
         
         # x = self.recon_trunk_light(x)
-        
-        for decoder, up, enc_skip in zip(self.decoders, self.ups, encs[::-1]):
+        i = 0
+        for decoder, up, fac_skip in zip(self.decoders, self.ups, facs[::-1]):
             x = up(x)
-            x = x + enc_skip
-            x = decoder(x)
+            if i == 2: # in the toppest decoder step
+                x = x + fac_skip
+                x = decoder(x)
+            else:
+                x = x + fac_skip
+                x = decoder(x)
+            i+=1
 
         x = self.ending(x)
         x = x + input
@@ -121,19 +139,29 @@ if __name__ == '__main__':
     img_channel = 3
     width = 32
 
+    # enc_blks = [1, 1, 1, 3]
+    # middle_blk_num = 3
+    # dec_blks = [2, 1, 1, 1]
+    
     enc_blks = [1, 2, 3]
-    middle_blk_num = 3
+    middle_blk_num_enc = 2
+    middle_blk_num_dec = 2
     dec_blks = [3, 1, 1]
-    residual_layers = 2
-    dilations = [1, 4]
+    residual_layers = None
+    dilations = [1, 4, 9]
+    extra_depth_wise = True
+    ksize = 5
     
     net = Network(img_channel=img_channel, 
                   width=width, 
-                  middle_blk_num=middle_blk_num,
+                  middle_blk_num_enc=middle_blk_num_enc,
+                  middle_blk_num_dec= middle_blk_num_dec,
                   enc_blk_nums=enc_blks, 
                   dec_blk_nums=dec_blks,
-                  dilations = dilations)
-
+                  dilations = dilations,
+                  extra_depth_wise = extra_depth_wise,
+                  ksize = ksize)
+    
     # NAF = NAFNet(img_channel=img_channel, width=width, middle_blk_num=middle_blk_num,
     #                   enc_blk_nums=enc_blks, dec_blk_nums=dec_blks)