Update archs/arch_util.py

archs/arch_util.py

@@ -10,128 +10,6 @@ except:
     from nafnet_utils.arch_util import LayerNorm2d
     from nafnet_utils.arch_model import SimpleGate
 
-'''
-https://github.com/wangchx67/FourLLIE.git
-'''
-
-def initialize_weights(net_l, scale=1):
-    if not isinstance(net_l, list):
-        net_l = [net_l]
-    for net in net_l:
-        for m in net.modules():
-            if isinstance(m, nn.Conv2d):
-                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
-                m.weight.data *= scale  # for residual block
-                if m.bias is not None:
-                    m.bias.data.zero_()
-            elif isinstance(m, nn.Linear):
-                init.kaiming_normal_(m.weight, a=0, mode='fan_in')
-                m.weight.data *= scale
-                if m.bias is not None:
-                    m.bias.data.zero_()
-            elif isinstance(m, nn.BatchNorm2d):
-                init.constant_(m.weight, 1)
-                init.constant_(m.bias.data, 0.0)
-
-
-def make_layer(block, n_layers):
-    layers = []
-    for _ in range(n_layers):
-        layers.append(block())
-    return nn.Sequential(*layers)
-
-
-class ResidualBlock_noBN(nn.Module):
-    '''Residual block w/o BN
-    ---Conv-ReLU-Conv-+-
-     |________________|
-    '''
-
-    def __init__(self, nf=64):
-        super(ResidualBlock_noBN, self).__init__()
-        self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
-        self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
-
-        # initialization
-        initialize_weights([self.conv1, self.conv2], 0.1)
-
-    def forward(self, x):
-        identity = x
-        out = F.relu(self.conv1(x), inplace=True)
-        out = self.conv2(out)
-        return identity + out
-
-class SpaBlock(nn.Module):
-    def __init__(self, nc):
-        super(SpaBlock, self).__init__()
-        self.block = nn.Sequential(
-            nn.Conv2d(nc,nc,3,1,1),
-            nn.LeakyReLU(0.1,inplace=True),
-            nn.Conv2d(nc, nc, 3, 1, 1),
-            nn.LeakyReLU(0.1, inplace=True))
-
-    def forward(self, x):
-        return x+self.block(x)
-
-class FreBlock(nn.Module):
-    def __init__(self, nc):
-        super(FreBlock, self).__init__()
-        self.fpre = nn.Conv2d(nc, nc, 1, 1, 0)
-        self.process1 = nn.Sequential(
-            nn.Conv2d(nc, nc, 1, 1, 0),
-            nn.LeakyReLU(0.1, inplace=True),
-            nn.Conv2d(nc, nc, 1, 1, 0))
-        self.process2 = nn.Sequential(
-            nn.Conv2d(nc, nc, 1, 1, 0),
-            nn.LeakyReLU(0.1, inplace=True),
-            nn.Conv2d(nc, nc, 1, 1, 0))
-
-    def forward(self, x):
-        _, _, H, W = x.shape
-        x_freq = torch.fft.rfft2(self.fpre(x), norm='backward')
-        mag = torch.abs(x_freq)
-        pha = torch.angle(x_freq)
-        mag = self.process1(mag)
-        pha = self.process2(pha)
-        real = mag * torch.cos(pha)
-        imag = mag * torch.sin(pha)
-        x_out = torch.complex(real, imag)
-        x_out = torch.fft.irfft2(x_out, s=(H, W), norm='backward')
-
-        return x_out+x
-
-class ProcessBlock(nn.Module):
-    def __init__(self, in_nc, spatial = True):
-        super(ProcessBlock,self).__init__()
-        self.spatial = spatial
-        self.spatial_process = SpaBlock(in_nc) if spatial else nn.Identity()
-        self.frequency_process = FreBlock(in_nc)
-        self.cat = nn.Conv2d(2*in_nc,in_nc,1,1,0) if spatial else nn.Conv2d(in_nc,in_nc,1,1,0)
-
-    def forward(self, x):
-        xori = x
-        x_freq = self.frequency_process(x)
-        x_spatial = self.spatial_process(x)
-        xcat = torch.cat([x_spatial,x_freq],1)
-        x_out = self.cat(xcat) if self.spatial else self.cat(x_freq)
-
-        return x_out+xori
-
-class Attention_Light(nn.Module):
-    
-    def __init__(self, img_channels = 3, width = 16, spatial = False):
-        super(Attention_Light, self).__init__()
-        self.block = nn.Sequential(
-                nn.Conv2d(in_channels = img_channels, out_channels = width//2, kernel_size = 1, padding = 0, stride = 1, groups = 1, bias = True),
-                ProcessBlock(in_nc = width //2, spatial = spatial),
-                nn.Conv2d(in_channels = width//2, out_channels = width, kernel_size = 1, padding = 0, stride = 1, groups = 1, bias = True),
-                ProcessBlock(in_nc = width, spatial = spatial),
-                nn.Conv2d(in_channels = width, out_channels = width, kernel_size = 1, padding = 0, stride = 1, groups = 1, bias = True),
-                ProcessBlock(in_nc=width, spatial = spatial),
-                nn.Sigmoid()
-                    )
-    def forward(self, input):
-        return self.block(input)
 
 class Branch(nn.Module):
     '''
@@ -223,7 +101,18 @@ if __name__ == '__main__':
 
     from ptflops import get_model_complexity_info
 
-    macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=True)
+    # macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=True)
+
+    # print('Values of EBlock:')
+    # print(macs, params)
 
+    channels = 128
+    resol = 32
+    ksize = 5
 
+    net = FAC(channels=channels, ksize=ksize)
+    inp_shape = (channels, resol, resol)
+    macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=True)
+    print('Values of FAC:')
     print(macs, params)
+