Update models/restormer_arch.py

models/restormer_arch.py

@@ -107,12 +107,11 @@ class Attention(nn.Module):
         self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
         
 
-
     def forward(self, x):
         b,c,h,w = x.shape
 
         qkv = self.qkv_dwconv(self.qkv(x))
-        q,k,v = qkv.chunk(3, dim=1)   
+        q,k,v = qkv.chunk(3, dim=1)    
         
         q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
         k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
@@ -200,8 +199,8 @@ class Restormer(nn.Module):
         heads = [1,2,4,8],
         ffn_expansion_factor = 2.66,
         bias = False,
-        LayerNorm_type = 'WithBias',   ## Other option 'BiasFree'
-        dual_pixel_task = True        ## True for dual-pixel defocus deblurring only. Also set inp_channels=6
+        LayerNorm_type = 'WithBias',    ## Other option 'BiasFree'
+        dual_pixel_task = True         ## True for dual-pixel defocus deblurring only. Also set inp_channels=6
     ):
 
         super(Restormer, self).__init__()
@@ -257,7 +256,7 @@ class Restormer(nn.Module):
         inp_enc_level4 = self.down3_4(out_enc_level3)        
         latent = self.latent(inp_enc_level4)
 
-                        
+                            
         inp_dec_level3 = self.up4_3(latent)
         inp_dec_level3 = torch.cat([inp_dec_level3, out_enc_level3], 1)
         inp_dec_level3 = self.reduce_chan_level3(inp_dec_level3)
@@ -274,7 +273,8 @@ class Restormer(nn.Module):
         
         out_dec_level1 = self.refinement(out_dec_level1)
 
-        out_dec_level1 = out_dec_level1 + self.skip_conv(inp_enc_level1)
+        if self.dual_pixel_task:
+             out_dec_level1 = out_dec_level1 + self.skip_conv(inp_enc_level1)
         out_dec_level1 = self.output(out_dec_level1)
 
         return out_dec_level1
@@ -283,6 +283,9 @@ class Restormer(nn.Module):
 
 if __name__ == '__main__':
     from torchtoolbox.tools import summary
+    # NOTE: The thop and torchtoolbox imports might require installation (pip install thop torchtoolbox)
+    # The summary function from torchtoolbox is used here.
+
     model = Restormer( 
         inp_channels=6, 
         out_channels=3, 
@@ -293,16 +296,31 @@ if __name__ == '__main__':
         heads = [1,2,4,8],
         ffn_expansion_factor = 2.66,
         bias = False,
-        LayerNorm_type = 'WithBias',   ## Other option 'BiasFree'
-        dual_pixel_task = True        ## True for dual-pixel defocus deblurring only. Also set inp_channels=6
+        LayerNorm_type = 'WithBias',    ## Other option 'BiasFree'
+        dual_pixel_task = True         ## True for dual-pixel defocus deblurring only. Also set inp_channels=6
     )
-    # model = Restormer(num_blocks=[4, 6, 6, 8], num_heads=[1, 2, 4, 8], channels=[48, 96, 192, 384], num_refinement=4, expansion_factor=2.66)
-    print(summary(model,torch.rand((1, 6, 256, 256))))
-
-    from thop import profile 
-    input = torch.rand((1, 6, 256, 256))
-    gflops,params = profile(model,inputs=(input,))
-    gflops = gflops*2 / 10**9
-    params = params / 10**6
-    print(gflops,'==============')
-    print(params,'==============')
+
+    print("--- Model Summary (Requires torchtoolbox) ---")
+    # print(summary(model,torch.rand((1, 6, 256, 256)))) # Uncomment if torchtoolbox is installed
+
+    print("\n--- Model Profiling (Requires thop) ---")
+    try:
+        from thop import profile 
+        input = torch.rand((1, 6, 256, 256))
+        gflops,params = profile(model,inputs=(input,), verbose=False)
+        gflops = gflops*2 / 10**9
+        params = params / 10**6
+        print(f"GFLOPS: {gflops:.4f}")
+        print(f"Params (M): {params:.4f}")
+    except ImportError:
+        print("Note: 'thop' library not found. Skipping GFLOPS/Params calculation.")
+    except Exception as e:
+        print(f"An error occurred during profiling: {e}")
+
+    # Example of a simple forward pass test
+    try:
+        input_tensor = torch.rand((1, 6, 256, 256))
+        output_tensor = model(input_tensor)
+        print(f"\nForward Pass Test: Input Shape {input_tensor.shape}, Output Shape {output_tensor.shape}")
+    except Exception as e:
+        print(f"\nForward Pass failed: {e}")