Update the model codes, including the previous inconsistencies.

models/backbones/build_backbone.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 from collections import OrderedDict
 from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights
-from models.backbones.pvt_v2 import pvt_v2_b2, pvt_v2_b5
+from models.backbones.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5
 from models.backbones.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l
 from config import Config
 

models/backbones/swin_v1.py

@@ -578,31 +578,6 @@ class SwinTransformer(nn.Module):
                 for param in m.parameters():
                     param.requires_grad = False
 
-    def init_weights(self, pretrained=None):
-        """Initialize the weights in backbone.
-
-        Args:
-            pretrained (str, optional): Path to pre-trained weights.
-                Defaults to None.
-        """
-
-        def _init_weights(m):
-            if isinstance(m, nn.Linear):
-                trunc_normal_(m.weight, std=.02)
-                if isinstance(m, nn.Linear) and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, nn.LayerNorm):
-                nn.init.constant_(m.bias, 0)
-                nn.init.constant_(m.weight, 1.0)
-
-        if isinstance(pretrained, str):
-            self.apply(_init_weights)
-            logger = get_root_logger()
-            load_checkpoint(self, pretrained, strict=False, logger=logger)
-        elif pretrained is None:
-            self.apply(_init_weights)
-        else:
-            raise TypeError('pretrained must be a str or None')
 
     def forward(self, x):
         """Forward function."""

models/{baseline.py → birefnet.py}

@@ -41,14 +41,6 @@ class BiRefNet(nn.Module):
             ])
 
         self.decoder = Decoder(channels)
-        
-        if self.config.locate_head:
-            self.locate_header = nn.ModuleList([
-                BasicDecBlk(channels[0], channels[-1]),
-                nn.Sequential(
-                    nn.Conv2d(channels[-1], 1, 1, 1, 0),
-                )
-            ])
 
         if self.config.ender:
             self.dec_end = nn.Sequential(
@@ -60,7 +52,7 @@ class BiRefNet(nn.Module):
         # refine patch-level segmentation
         if self.config.refine:
             if self.config.refine == 'itself':
-                self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3)
+                self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN')
             else:
                 self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1'))
 
@@ -105,20 +97,6 @@ class BiRefNet(nn.Module):
             )
         return (x1, x2, x3, x4), class_preds
 
-    # def forward_loc(self, x):
-    #     ########## Encoder ##########
-    #     (x1, x2, x3, x4), class_preds = self.forward_enc(x)
-    #     if self.config.squeeze_block:
-    #         x4 = self.squeeze_module(x4)
-    #     if self.config.locate_head:
-    #         locate_preds = self.locate_header[1](
-    #             F.interpolate(
-    #                 self.locate_header[0](
-    #                     F.interpolate(x4, size=x2.shape[2:], mode='bilinear', align_corners=True)
-    #                 ), size=x.shape[2:], mode='bilinear', align_corners=True
-    #             )
-    #         )
-
     def forward_ori(self, x):
         ########## Encoder ##########
         (x1, x2, x3, x4), class_preds = self.forward_enc(x)
@@ -131,22 +109,22 @@ class BiRefNet(nn.Module):
         scaled_preds = self.decoder(features)
         return scaled_preds, class_preds
 
-    def forward_ref(self, x, pred):
-        # refine patch-level segmentation
-        if pred.shape[2:] != x.shape[2:]:
-            pred = F.interpolate(pred, size=x.shape[2:], mode='bilinear', align_corners=True)
-        # pred = pred.sigmoid()
-        if self.config.refine == 'itself':
-            x = self.stem_layer(torch.cat([x, pred], dim=1))
-            scaled_preds, class_preds = self.forward_ori(x)
-        else:
-            scaled_preds = self.refiner([x, pred])
-            class_preds = None
-        return scaled_preds, class_preds
+    # def forward_ref(self, x, pred):
+    #     # refine patch-level segmentation
+    #     if pred.shape[2:] != x.shape[2:]:
+    #         pred = F.interpolate(pred, size=x.shape[2:], mode='bilinear', align_corners=True)
+    #     # pred = pred.sigmoid()
+    #     if self.config.refine == 'itself':
+    #         x = self.stem_layer(torch.cat([x, pred], dim=1))
+    #         scaled_preds, class_preds = self.forward_ori(x)
+    #     else:
+    #         scaled_preds = self.refiner([x, pred])
+    #         class_preds = None
+    #     return scaled_preds, class_preds
 
-    def forward_ref_end(self, x):
-        # remove the grids of concatenated preds
-        return self.dec_end(x) if self.config.ender else x
+    # def forward_ref_end(self, x):
+    #     # remove the grids of concatenated preds
+    #     return self.dec_end(x) if self.config.ender else x
 
 
     # def forward(self, x):
@@ -181,6 +159,7 @@ class Decoder(nn.Module):
             DBlock = SimpleConvs
             ic = 64
             ipt_cha_opt = 1
+            self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic)
             self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic)
             self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic)
             self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic)
@@ -188,7 +167,7 @@ class Decoder(nn.Module):
         else:
             self.split = None
 
-        self.decoder_block4 = DecoderBlock(channels[0], channels[1])
+        self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1])
         self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2])
         self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3])
         self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2)
@@ -205,15 +184,15 @@ class Decoder(nn.Module):
 
             if self.config.out_ref:
                 _N = 16
-                # self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N), nn.ReLU(inplace=True))
-                self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N), nn.ReLU(inplace=True))
-                self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N), nn.ReLU(inplace=True))
+                self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
+                self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
+                self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
 
-                # self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+                self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 
-                # self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+                self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
 
@@ -238,14 +217,31 @@ class Decoder(nn.Module):
         else:
             x, x1, x2, x3, x4 = features
         outs = []
+
+        if self.config.dec_ipt:
+            patches_batch = self.get_patches_batch(x, x4) if self.split else x
+            x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1)
         p4 = self.decoder_block4(x4)
         m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision else None
+        if self.config.out_ref:
+            p4_gdt = self.gdt_convs_4(p4)
+            if self.training:
+                # >> GT:
+                m4_dia = m4
+                gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
+                outs_gdt_label.append(gdt_label_main_4)
+                # >> Pred:
+                gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt)
+                outs_gdt_pred.append(gdt_pred_4)
+            gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid()
+            # >> Finally:
+            p4 = p4 * gdt_attn_4
         _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True)
         _p3 = _p4 + self.lateral_block4(x3)
+
         if self.config.dec_ipt:
             patches_batch = self.get_patches_batch(x, _p3) if self.split else x
             _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1)
-
         p3 = self.decoder_block3(_p3)
         m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision else None
         if self.config.out_ref:
@@ -268,10 +264,10 @@ class Decoder(nn.Module):
             p3 = p3 * gdt_attn_3
         _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True)
         _p2 = _p3 + self.lateral_block3(x2)
+
         if self.config.dec_ipt:
             patches_batch = self.get_patches_batch(x, _p2) if self.split else x
             _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1)
-
         p2 = self.decoder_block2(_p2)
         m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision else None
         if self.config.out_ref:
@@ -289,12 +285,13 @@ class Decoder(nn.Module):
             p2 = p2 * gdt_attn_2
         _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True)
         _p1 = _p2 + self.lateral_block2(x1)
+
         if self.config.dec_ipt:
             patches_batch = self.get_patches_batch(x, _p1) if self.split else x
             _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1)
-
         _p1 = self.decoder_block1(_p1)
         _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True)
+
         if self.config.dec_ipt:
             patches_batch = self.get_patches_batch(x, _p1) if self.split else x
             _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1)

models/modules/aspp.py

@@ -8,56 +8,12 @@ from config import Config
 config = Config()
 
 
-class ASPPComplex(nn.Module):
-    def __init__(self, in_channels=64, out_channels=None, output_stride=16):
-        super(ASPPComplex, self).__init__()
-        self.down_scale = 1
-        if out_channels is None:
-            out_channels = in_channels
-        self.in_channelster = 256 // self.down_scale
-        if output_stride == 16:
-            dilations = [1, 6, 12, 18]
-        elif output_stride == 8:
-            dilations = [1, 12, 24, 36]
-        else:
-            raise NotImplementedError
-
-        self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0])
-        self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1])
-        self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2])
-        self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3])
-
-        self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
-                                             nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False),
-                                             nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(),
-                                             nn.ReLU(inplace=True))
-        self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False)
-        self.bn1 = nn.BatchNorm2d(out_channels)
-        self.relu = nn.ReLU(inplace=True)
-        self.dropout = nn.Dropout(0.5)
-
-    def forward(self, x):
-        x1 = self.aspp1(x)
-        x2 = self.aspp2(x)
-        x3 = self.aspp3(x)
-        x4 = self.aspp4(x)
-        x5 = self.global_avg_pool(x)
-        x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True)
-        x = torch.cat((x1, x2, x3, x4, x5), dim=1)
-
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.relu(x)
-
-        return self.dropout(x)
-
-
 class _ASPPModule(nn.Module):
     def __init__(self, in_channels, planes, kernel_size, padding, dilation):
         super(_ASPPModule, self).__init__()
         self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size,
                                             stride=1, padding=padding, dilation=dilation, bias=False)
-        self.bn = nn.BatchNorm2d(planes)
+        self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity()
         self.relu = nn.ReLU(inplace=True)
 
     def forward(self, x):
@@ -66,6 +22,7 @@ class _ASPPModule(nn.Module):
 
         return self.relu(x)
 
+
 class ASPP(nn.Module):
     def __init__(self, in_channels=64, out_channels=None, output_stride=16):
         super(ASPP, self).__init__()
@@ -90,7 +47,7 @@ class ASPP(nn.Module):
                                              nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(),
                                              nn.ReLU(inplace=True))
         self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False)
-        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
         self.relu = nn.ReLU(inplace=True)
         self.dropout = nn.Dropout(0.5)
 
@@ -116,7 +73,7 @@ class _ASPPModuleDeformable(nn.Module):
         super(_ASPPModuleDeformable, self).__init__()
         self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size,
                                             stride=1, padding=padding, bias=False)
-        self.bn = nn.BatchNorm2d(planes)
+        self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity()
         self.relu = nn.ReLU(inplace=True)
 
     def forward(self, x):
@@ -127,7 +84,7 @@ class _ASPPModuleDeformable(nn.Module):
 
 
 class ASPPDeformable(nn.Module):
-    def __init__(self, in_channels, out_channels=None, num_parallel_block=1):
+    def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]):
         super(ASPPDeformable, self).__init__()
         self.down_scale = 1
         if out_channels is None:
@@ -136,7 +93,7 @@ class ASPPDeformable(nn.Module):
 
         self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0)
         self.aspp_deforms = nn.ModuleList([
-            _ASPPModuleDeformable(in_channels, self.in_channelster, 3, padding=1) for _ in range(num_parallel_block)
+            _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes
         ])
 
         self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
@@ -144,7 +101,7 @@ class ASPPDeformable(nn.Module):
                                              nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(),
                                              nn.ReLU(inplace=True))
         self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False)
-        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
         self.relu = nn.ReLU(inplace=True)
         self.dropout = nn.Dropout(0.5)
 

models/modules/decoder_blocks.py

@@ -19,8 +19,8 @@ class BasicDecBlk(nn.Module):
         elif config.dec_att == 'ASPPDeformable':
             self.dec_att = ASPPDeformable(in_channels=inter_channels)
         self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
-        self.bn_in = nn.BatchNorm2d(inter_channels)
-        self.bn_out = nn.BatchNorm2d(out_channels)
+        self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity()
+        self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
 
     def forward(self, x):
         x = self.conv_in(x)
@@ -41,7 +41,7 @@ class ResBlk(nn.Module):
         inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64
 
         self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1)
-        self.bn_in = nn.BatchNorm2d(inter_channels)
+        self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity()
         self.relu_in = nn.ReLU(inplace=True)
 
         if config.dec_att == 'ASPP':
@@ -50,7 +50,7 @@ class ResBlk(nn.Module):
             self.dec_att = ASPPDeformable(in_channels=inter_channels)
 
         self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
-        self.bn_out = nn.BatchNorm2d(out_channels)
+        self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
         
         self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
 

models/refinement/refiner.py

@@ -65,7 +65,7 @@ class Refiner(nn.Module):
         super(Refiner, self).__init__()
         self.config = Config()
         self.epoch = 1
-        self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3)
+        self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN')
         self.bb = build_backbone(self.config.bb)
 
         lateral_channels_in_collection = {