chore: Update SCHP model checkpoint loading logic

.gitignore

@@ -0,0 +1,2 @@
+playground.py
+__pycache__

model/SCHP/__init__.py

@@ -81,12 +81,27 @@ class SCHP:
 
 
     def load_ckpt(self, ckpt_path):
+        rename_map = {
+            "decoder.conv3.2.weight": "decoder.conv3.3.weight",
+            "decoder.conv3.3.weight": "decoder.conv3.4.weight",
+            "decoder.conv3.3.bias": "decoder.conv3.4.bias",
+            "decoder.conv3.3.running_mean": "decoder.conv3.4.running_mean",
+            "decoder.conv3.3.running_var": "decoder.conv3.4.running_var",
+            "fushion.3.weight": "fushion.4.weight",
+            "fushion.3.bias": "fushion.4.bias",
+        }
         state_dict = torch.load(ckpt_path, map_location='cpu')['state_dict']
         new_state_dict = OrderedDict()
         for k, v in state_dict.items():
             name = k[7:]  # remove `module.`
             new_state_dict[name] = v
-        self.model.load_state_dict(new_state_dict)
+        new_state_dict_ = OrderedDict()
+        for k, v in list(new_state_dict.items()):
+            if k in rename_map:
+                new_state_dict_[rename_map[k]] = v
+            else:
+                new_state_dict_[k] = v
+        self.model.load_state_dict(new_state_dict_, strict=False)
 
     def _box2cs(self, box):
         x, y, w, h = box[:4]
@@ -148,7 +163,8 @@ class SCHP:
             meta_list = [meta]
                 
         output = self.model(image)
-        upsample_outputs = self.upsample(output[0][-1])
+        # upsample_outputs = self.upsample(output[0][-1])
+        upsample_outputs = self.upsample(output)
         upsample_outputs = upsample_outputs.permute(0, 2, 3, 1)  # BCHW -> BHWC
 
         output_img_list = []

model/SCHP/networks/AugmentCE2P.py

@@ -11,19 +11,13 @@
              LICENSE file in the root directory of this source tree.
 """
 
-import functools
-
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
-# Note here we adopt the InplaceABNSync implementation from https://github.com/mapillary/inplace_abn
-# By default, the InplaceABNSync module contains a BatchNorm Layer and a LeakyReLu layer
-from inplace_abn import InPlaceABNSync
 
-BatchNorm2d = functools.partial(InPlaceABNSync, activation='identity')
+from torch.nn import BatchNorm2d, LeakyReLU
 
 affine_par = True
-
 pretrained_settings = {
     'resnet101': {
         'imagenet': {
@@ -99,14 +93,20 @@ class PSPModule(nn.Module):
         self.bottleneck = nn.Sequential(
             nn.Conv2d(features + len(sizes) * out_features, out_features, kernel_size=3, padding=1, dilation=1,
                       bias=False),
-            InPlaceABNSync(out_features),
+            BatchNorm2d(out_features),
+            LeakyReLU(),
         )
 
     def _make_stage(self, features, out_features, size):
         prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
         conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
-        bn = InPlaceABNSync(out_features)
-        return nn.Sequential(prior, conv, bn)
+        return nn.Sequential(
+            prior, 
+            conv, 
+            # bn
+            BatchNorm2d(out_features),
+            LeakyReLU(),
+        )
 
     def forward(self, feats):
         h, w = feats.size(2), feats.size(3)
@@ -128,23 +128,35 @@ class ASPPModule(nn.Module):
         self.conv1 = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
                                    nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1,
                                              bias=False),
-                                   InPlaceABNSync(inner_features))
+                                #    InPlaceABNSync(inner_features)
+                                    BatchNorm2d(inner_features),
+                                    LeakyReLU(),
+                                )
         self.conv2 = nn.Sequential(
             nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(inner_features))
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            )
         self.conv3 = nn.Sequential(
             nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=False),
-            InPlaceABNSync(inner_features))
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            )
         self.conv4 = nn.Sequential(
             nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[1], dilation=dilations[1], bias=False),
-            InPlaceABNSync(inner_features))
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            )
         self.conv5 = nn.Sequential(
             nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[2], dilation=dilations[2], bias=False),
-            InPlaceABNSync(inner_features))
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
+            )
 
         self.bottleneck = nn.Sequential(
             nn.Conv2d(inner_features * 5, out_features, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(out_features),
+            BatchNorm2d(inner_features),
+            LeakyReLU(),
             nn.Dropout2d(0.1)
         )
 
@@ -173,24 +185,27 @@ class Edge_Module(nn.Module):
 
         self.conv1 = nn.Sequential(
             nn.Conv2d(in_fea[0], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(mid_fea)
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
         )
         self.conv2 = nn.Sequential(
             nn.Conv2d(in_fea[1], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(mid_fea)
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
         )
         self.conv3 = nn.Sequential(
             nn.Conv2d(in_fea[2], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(mid_fea)
+            BatchNorm2d(mid_fea),
+            LeakyReLU(),
         )
         self.conv4 = nn.Conv2d(mid_fea, out_fea, kernel_size=3, padding=1, dilation=1, bias=True)
-        self.conv5 = nn.Conv2d(out_fea * 3, out_fea, kernel_size=1, padding=0, dilation=1, bias=True)
+        # self.conv5 = nn.Conv2d(out_fea * 3, out_fea, kernel_size=1, padding=0, dilation=1, bias=True)
 
     def forward(self, x1, x2, x3):
         _, _, h, w = x1.size()
 
         edge1_fea = self.conv1(x1)
-        edge1 = self.conv4(edge1_fea)
+        # edge1 = self.conv4(edge1_fea)
         edge2_fea = self.conv2(x2)
         edge2 = self.conv4(edge2_fea)
         edge3_fea = self.conv3(x3)
@@ -201,11 +216,12 @@ class Edge_Module(nn.Module):
         edge2 = F.interpolate(edge2, size=(h, w), mode='bilinear', align_corners=True)
         edge3 = F.interpolate(edge3, size=(h, w), mode='bilinear', align_corners=True)
 
-        edge = torch.cat([edge1, edge2, edge3], dim=1)
+        # edge = torch.cat([edge1, edge2, edge3], dim=1)
         edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
-        edge = self.conv5(edge)
+        # edge = self.conv5(edge)
 
-        return edge, edge_fea
+        # return edge, edge_fea
+        return edge_fea
 
 
 class Decoder_Module(nn.Module):
@@ -217,20 +233,24 @@ class Decoder_Module(nn.Module):
         super(Decoder_Module, self).__init__()
         self.conv1 = nn.Sequential(
             nn.Conv2d(512, 256, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(256)
+            BatchNorm2d(256),
+            LeakyReLU(),
         )
         self.conv2 = nn.Sequential(
             nn.Conv2d(256, 48, kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(48)
+            BatchNorm2d(48),
+            LeakyReLU(),
         )
         self.conv3 = nn.Sequential(
             nn.Conv2d(304, 256, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(256),
+            BatchNorm2d(256),
+            LeakyReLU(),
             nn.Conv2d(256, 256, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(256)
+            BatchNorm2d(256),
+            LeakyReLU(),
         )
 
-        self.conv4 = nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
+        # self.conv4 = nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
 
     def forward(self, xt, xl):
         _, _, h, w = xl.size()
@@ -238,8 +258,9 @@ class Decoder_Module(nn.Module):
         xl = self.conv2(xl)
         x = torch.cat([xt, xl], dim=1)
         x = self.conv3(x)
-        seg = self.conv4(x)
-        return seg, x
+        # seg = self.conv4(x)
+        # return seg, x
+        return x
 
 
 class ResNet(nn.Module):
@@ -270,7 +291,8 @@ class ResNet(nn.Module):
 
         self.fushion = nn.Sequential(
             nn.Conv2d(1024, 256, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(256),
+            BatchNorm2d(256),
+            LeakyReLU(),
             nn.Dropout2d(0.1),
             nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
         )
@@ -304,13 +326,16 @@ class ResNet(nn.Module):
         x4 = self.layer3(x3)
         x5 = self.layer4(x4)
         x = self.context_encoding(x5)
-        parsing_result, parsing_fea = self.decoder(x, x2)
+        # parsing_result, parsing_fea = self.decoder(x, x2)
+        parsing_fea = self.decoder(x, x2)
         # Edge Branch
-        edge_result, edge_fea = self.edge(x2, x3, x4)
+        # edge_result, edge_fea = self.edge(x2, x3, x4)
+        edge_fea = self.edge(x2, x3, x4)
         # Fusion Branch
         x = torch.cat([parsing_fea, edge_fea], dim=1)
         fusion_result = self.fushion(x)
-        return [[parsing_result, fusion_result], [edge_result]]
+        # return [[parsing_result, fusion_result], [edge_result]]
+        return fusion_result
 
 
 def initialize_pretrained_model(model, settings, pretrained='./models/resnet101-imagenet.pth'):