Update networks/decoder.py

networks/decoder.py

@@ -5,9 +5,11 @@ import torch.nn as nn
 from torch import Tensor
 import torch.nn.functional as F
 
+from mmcv.cnn import ConvModule
+from .height_head import resize
+
 class PPM(nn.ModuleList):
     """Pooling Pyramid Module used in PSPNet.
-
     Args:
         pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid
             Module.
@@ -155,3 +157,24 @@ class Decoder(nn.Module):
             return torch.cat(psp_outs, dim=1)
         else:
             return x
+
+if __name__ == '__main__':
+    # model = Decoder(in_channel=320)
+    # input_data = torch.randn(1, 320, 32, 32)
+    # res = [torch.randn(1, 512, 32, 32), torch.randn(1, 256, 64, 64), torch.randn(1, 128, 128, 128)]
+    # output = model(input_data, res)
+    # print(output.shape)
+
+    model = Decoder(in_channel=320, short_cut_channels=None)
+    # input_data = torch.randn(1, 320, 32, 32)
+    # output = model(input_data, None)
+    # flops, params = get_model_complexity_info(model, (320, 32, 32))
+    # print(f"参数量: {params}")
+    # print(f"计算量: {flops}")
+    # print("-" * 30)
+    # print(output.shape)
+
+    # model = Decoder(in_channel=320, short_cut_channels=None, psp_channel=-1)
+    # input_data = torch.randn(2, 320, 32, 32)
+    # output = model(input_data, None)
+    # print(output.shape)