Update model/SUNet_detail.py

model/SUNet_detail.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 import torch.utils.checkpoint as checkpoint
 from einops import rearrange
 from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-from thop import profile
+
 
 class Mlp(nn.Module):
     def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
@@ -763,26 +763,3 @@ class SUNet(nn.Module):
         flops += self.num_features * self.out_chans
         return flops
 
-
-if __name__ == '__main__':
-    from utils.model_utils import network_parameters
-
-    height = 256
-    width = 256
-    x = torch.randn((1, 3, height, width))  # .cuda()
-    model = SUNet(img_size=256, patch_size=4, in_chans=3, out_chans=3,
-                  embed_dim=96, depths=[8, 8, 8, 8],
-                  num_heads=[8, 8, 8, 8],
-                  window_size=8, mlp_ratio=4., qkv_bias=True, qk_scale=2,
-                  drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
-                  norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
-                  use_checkpoint=False, final_upsample="Dual up-sample")  # .cuda()
-    # print(model)
-    print('input image size: (%d, %d)' % (height, width))
-    print('FLOPs: %.4f G' % (model.flops() / 1e9))
-    print('model parameters: ', network_parameters(model))
-    # x = model(x)
-    print('output image size: ', x.shape)
-    flops, params = profile(model, (x,))
-    print(flops)
-    print(params)