add upernet for ade20k

Author: Tete Xiao

README.md

@@ -137,7 +137,7 @@ usage: train.py [-h] [--id ID] [--arch_encoder ARCH_ENCODER]
 
 
 ## Evaluation
-1. Evaluate a trained network on the validation set. Add ```--visualize``` option to output visualizations shown in teaser.
+1. Evaluate a trained network on the validation set. Add ```--visualize``` option to output visualizations as shown in teaser.
 ```bash
 python3 eval.py --id MODEL_ID --suffix SUFFIX
 ```

dataset.py

@@ -164,8 +164,6 @@ def __init__(self, odgt, opt, max_sample=-1):
         self.imgMaxSize = opt.imgMaxSize
         # max down sampling rate of network to avoid rounding during conv or pooling
         self.padding_constant = opt.padding_constant
-        # down sampling rate of segm labe
-        self.segm_downsampling_rate = opt.segm_downsampling_rate 
 
         # mean and std
         self.img_transform = transforms.Compose([

eval.py

@@ -165,8 +165,6 @@ def main(args):
                         help='maximum input image size of long edge')
     parser.add_argument('--padding_constant', default=8, type=int,
                         help='maxmimum downsampling rate of the network')
-    parser.add_argument('--segm_downsampling_rate', default=8, type=int,
-                        help='downsampling rate of the segmentation label')
 
     # Misc arguments
     parser.add_argument('--ckpt', default='./ckpt',

models.py

@@ -31,7 +31,7 @@ def forward(self, feed_dict, *, segSize=None):
             if self.deep_sup_scale is not None: # use deep supervision technique
                 (pred, pred_deepsup) = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True))
             else:
-                pred = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=False))
+                pred = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True))
 
             loss = self.crit(pred, feed_dict['seg_label'])
             if self.deep_sup_scale is not None:
@@ -41,7 +41,7 @@ def forward(self, feed_dict, *, segSize=None):
             acc = self.pixel_acc(pred, feed_dict['seg_label'])
             return loss, acc
         else: # inference
-            pred = self.decoder(self.encoder(feed_dict['img_data']), segSize=segSize)
+            pred = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True), segSize=segSize)
             return pred
 
 
@@ -94,10 +94,17 @@ def build_encoder(self, arch='resnet50_dilated8', fc_dim=512, weights=''):
             orig_resnet = resnet.__dict__['resnet50'](pretrained=pretrained)
             net_encoder = ResnetDilated(orig_resnet,
                                         dilate_scale=16)
+        elif arch == 'resnet101':
+            orig_resnet = resnet.__dict__['resnet101'](pretrained=pretrained)
+            net_encoder = Resnet(orig_resnet)
         elif arch == 'resnet101_dilated8':
             orig_resnet = resnet.__dict__['resnet101'](pretrained=pretrained)
             net_encoder = ResnetDilated(orig_resnet,
                                         dilate_scale=8)
+        elif arch == 'resnet101_dilated16':
+            orig_resnet = resnet.__dict__['resnet101'](pretrained=pretrained)
+            net_encoder = ResnetDilated(orig_resnet,
+                                        dilate_scale=16)
         else:
             raise Exception('Architecture undefined!')
 
@@ -131,6 +138,18 @@ def build_decoder(self, arch='psp_bilinear_deepsup',
                 num_class=num_class,
                 fc_dim=fc_dim,
                 use_softmax=use_softmax)
+        elif arch == 'upernet_lite':
+            net_decoder = UPerNet(
+                num_class=num_class,
+                fc_dim=fc_dim,
+                use_softmax=use_softmax,
+                fpn_dim=256)
+        elif arch == 'upernet':
+            net_decoder = UPerNet(
+                num_class=num_class,
+                fc_dim=fc_dim,
+                use_softmax=use_softmax,
+                fpn_dim=512)
         else:
             raise Exception('Architecture undefined!')
 
@@ -162,17 +181,22 @@ def __init__(self, orig_resnet):
         self.layer3 = orig_resnet.layer3
         self.layer4 = orig_resnet.layer4
 
-    def forward(self, x):
+    def forward(self, x, return_feature_maps=False):
+        conv_out = []
+
         x = self.relu1(self.bn1(self.conv1(x)))
         x = self.relu2(self.bn2(self.conv2(x)))
         x = self.relu3(self.bn3(self.conv3(x)))
         x = self.maxpool(x)
 
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        return x
+        x = self.layer1(x); conv_out.append(x);
+        x = self.layer2(x); conv_out.append(x);
+        x = self.layer3(x); conv_out.append(x);
+        x = self.layer4(x); conv_out.append(x);
+
+        if return_feature_maps:
+            return conv_out
+        return [x]
 
 
 class ResnetDilated(nn.Module):
@@ -404,3 +428,90 @@ def forward(self, conv_out, segSize=None):
         _ = nn.functional.log_softmax(_, dim=1)
 
         return (x, _)
+
+
+# upernet
+class UPerNet(nn.Module):
+    def __init__(self, num_class=150, fc_dim=4096,
+                 use_softmax=False, pool_scales=(1, 2, 3, 6),
+                 fpn_inplanes=(256,512,1024,2048), fpn_dim=256):
+        super(UPerNet, self).__init__()
+        self.use_softmax = use_softmax
+
+        # PPM Module
+        self.psp = []
+        for scale in pool_scales:
+            self.psp.append(nn.Sequential(
+                nn.AdaptiveAvgPool2d(scale),
+                nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
+                SynchronizedBatchNorm2d(512),
+                nn.ReLU(inplace=True)
+            ))
+        self.psp = nn.ModuleList(self.psp)
+        self.psp_conv = conv3x3_bn_relu(fc_dim + len(pool_scales)*512, fpn_dim, 1)
+
+        # FPN Module
+        self.fpn_in = []
+        for fpn_inplane in fpn_inplanes[:-1]: # skip the top layer
+            self.fpn_in.append(nn.Sequential(
+                nn.Conv2d(fpn_inplane, fpn_dim, kernel_size=1, bias=False),
+                SynchronizedBatchNorm2d(fpn_dim),
+                nn.ReLU(inplace=True)
+            ))
+        self.fpn_in = nn.ModuleList(self.fpn_in)
+
+        self.fpn_out = []
+        for i in range(len(fpn_inplanes) - 1): # skip the top layer
+            self.fpn_out.append(nn.Sequential(
+                conv3x3_bn_relu(fpn_dim, fpn_dim, 1),
+                conv3x3_bn_relu(fpn_dim, fpn_dim // 2, 1)
+            ))
+        self.fpn_out = nn.ModuleList(self.fpn_out)
+
+        self.conv_last = nn.Sequential(
+            conv3x3_bn_relu(fpn_dim + (len(fpn_inplanes)-1) * fpn_dim // 2, fpn_dim, 1),
+            nn.Conv2d(fpn_dim, num_class, kernel_size=1)
+        )
+
+    def forward(self, conv_out, segSize=None):
+        conv5 = conv_out[-1]
+
+        input_size = conv5.size()
+        psp_out = [conv5]
+        for pool_scale in self.psp:
+            psp_out.append(nn.functional.upsample(
+                pool_scale(conv5),
+                (input_size[2], input_size[3]),
+                mode='bilinear'))
+        psp_out = torch.cat(psp_out, 1)
+        f = self.psp_conv(psp_out)
+
+        fpn_feature_list = [f]
+        for i in reversed(range(len(conv_out) - 1)):
+            conv_x = conv_out[i]
+            conv_x = self.fpn_in[i](conv_x) # lateral branch
+
+            f = nn.functional.upsample(f, size=conv_x.size()[2:], mode='bilinear') # top-down branch
+            f = conv_x + f
+
+            fpn_feature_list.append(self.fpn_out[i](f))
+
+        fpn_feature_list.reverse() # [P2 - P5]
+        output_size = fpn_feature_list[0].size()[2:]
+        fusion_list = [fpn_feature_list[0]]
+        for i in range(1, len(fpn_feature_list)):
+            fusion_list.append(nn.functional.upsample(
+                fpn_feature_list[i],
+                output_size,
+                mode='bilinear'))
+        fusion_out = torch.cat(fusion_list, 1)
+        x = self.conv_last(fusion_out)
+
+        if self.use_softmax:  # is True during inference
+            x = nn.functional.upsample(x, size=segSize, mode='bilinear')
+            x = nn.functional.softmax(x, dim=1)
+            return x
+
+        x = nn.functional.log_softmax(x, dim=1)
+
+        return x