From 28f53173ae965650c28247daa02cfc93ee8ed236 Mon Sep 17 00:00:00 2001
From: liminn <11172918@qq.com>
Date: Wed, 13 Nov 2019 18:14:34 +0800
Subject: [PATCH] update icnet

---
 core/models/icnet.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/core/models/icnet.py b/core/models/icnet.py
index 62f0531a7..94d03444f 100644
--- a/core/models/icnet.py
+++ b/core/models/icnet.py
@@ -20,6 +20,8 @@ def __init__(self, nclass, backbone='resnet50', aux=False, jpu=False, pretrained
             _ConvBNReLU(32, 64, 3, 2, **kwargs)
         )
 
+        self.ppm = PyramidPoolingModule()
+
         self.head = _ICHead(nclass, **kwargs)
 
         self.__setattr__('exclusive', ['conv_sub1', 'head'])
@@ -35,16 +37,31 @@ def forward(self, x):
         # sub 4
         x_sub4 = F.interpolate(x, scale_factor=0.25, mode='bilinear', align_corners=True)
         _, _, _, x_sub4 = self.base_forward(x_sub4)
-
+        # add PyramidPoolingModule
+        x_sub4 = self.ppm(x_sub4)   
         outputs = self.head(x_sub1, x_sub2, x_sub4)
 
         return tuple(outputs)
 
+class PyramidPoolingModule(nn.Module):
+    def __init__(self, pyramids=[1,2,3,6]):
+        super(PyramidPoolingModule, self).__init__()
+        self.pyramids = pyramids
+
+    def forward(self, input):
+        feat = input
+        height, width = input.shape[2:]
+        for bin_size in self.pyramids:
+            x = F.adaptive_avg_pool2d(input, output_size=bin_size)
+            x = F.interpolate(x, size=(height, width), mode='bilinear', align_corners=True)
+            feat  = feat + x
+        return feat
 
 class _ICHead(nn.Module):
     def __init__(self, nclass, norm_layer=nn.BatchNorm2d, **kwargs):
         super(_ICHead, self).__init__()
-        self.cff_12 = CascadeFeatureFusion(512, 64, 128, nclass, norm_layer, **kwargs)
+        #self.cff_12 = CascadeFeatureFusion(512, 64, 128, nclass, norm_layer, **kwargs)
+        self.cff_12 = CascadeFeatureFusion(128, 64, 128, nclass, norm_layer, **kwargs)  
         self.cff_24 = CascadeFeatureFusion(2048, 512, 128, nclass, norm_layer, **kwargs)
 
         self.conv_cls = nn.Conv2d(128, nclass, 1, bias=False)
@@ -53,9 +70,10 @@ def forward(self, x_sub1, x_sub2, x_sub4):
         outputs = list()
         x_cff_24, x_24_cls = self.cff_24(x_sub4, x_sub2)
         outputs.append(x_24_cls)
-        x_cff_12, x_12_cls = self.cff_12(x_sub2, x_sub1)
+        #x_cff_12, x_12_cls = self.cff_12(x_sub2, x_sub1)
+        x_cff_12, x_12_cls = self.cff_12(x_cff_24, x_sub1)  
         outputs.append(x_12_cls)
-
+        
         up_x2 = F.interpolate(x_cff_12, scale_factor=2, mode='bilinear', align_corners=True)
         up_x2 = self.conv_cls(up_x2)
         outputs.append(up_x2)