update comments

Author: CharlesPikachu

ssseg/modules/models/backbones/beit.py

@@ -95,9 +95,7 @@ def forward(self, x):
         if self.relative_position_bias_table is not None:
             Wh = self.window_size[0]
             Ww = self.window_size[1]
-            relative_position_bias = self.relative_position_bias_table[
-                self.relative_position_index.view(-1)
-            ].view(Wh * Ww + 1, Wh * Ww + 1, -1)
+            relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(Wh * Ww + 1, Wh * Ww + 1, -1)
             relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
             attn = attn + relative_position_bias.unsqueeze(0)
         attn = attn.softmax(dim=-1)
@@ -122,13 +120,8 @@ def __init__(self, embed_dims, num_heads, feedforward_channels, attn_drop_rate=0
         self.gamma_1 = nn.Parameter(init_values * torch.ones((embed_dims)), requires_grad=True)
         self.gamma_2 = nn.Parameter(init_values * torch.ones((embed_dims)), requires_grad=True)
         attn_cfg.update(dict(
-            window_size=window_size,
-            qk_scale=None,
-            embed_dims=embed_dims,
-            num_heads=num_heads,
-            attn_drop_rate=attn_drop_rate,
-            proj_drop_rate=0.,
-            bias=bias,
+            window_size=window_size, qk_scale=None, embed_dims=embed_dims, num_heads=num_heads,
+            attn_drop_rate=attn_drop_rate, proj_drop_rate=0., bias=bias,
         ))
         self.attn = BEiTAttention(**attn_cfg)
     '''forward'''

ssseg/modules/models/backbones/bisenetv1.py

@@ -17,7 +17,7 @@
 AUTO_ASSERT_STRUCTURE_TYPES = {}
 
 
-'''Spatial Path to preserve the spatial size of the original input image and encode affluent spatial information'''
+'''SpatialPath'''
 class SpatialPath(nn.Module):
     def __init__(self, in_channels=3, num_channels_list=(64, 64, 64, 128), norm_cfg=None, act_cfg=None):
         super(SpatialPath, self).__init__()
@@ -53,7 +53,7 @@ def forward(self, x):
         return x
 
 
-'''Attention Refinement Module (ARM) to refine the features of each stage'''
+'''AttentionRefinementModule'''
 class AttentionRefinementModule(nn.Module):
     def __init__(self, in_channels, out_channels, norm_cfg=None, act_cfg=None):
         super(AttentionRefinementModule, self).__init__()
@@ -76,7 +76,7 @@ def forward(self, x):
         return x_out
 
 
-'''Context Path to provide sufficient receptive field'''
+'''ContextPath'''
 class ContextPath(nn.Module):
     def __init__(self, backbone_cfg, context_channels_list=(128, 256, 512), norm_cfg=None, act_cfg=None):
         super(ContextPath, self).__init__()
@@ -125,7 +125,7 @@ def buildbackbone(self, cfg):
         return supported_backbones[backbone_type](**cfg)
 
 
-'''Feature Fusion Module to fuse low level output feature of Spatial Path and high level output feature of Context Path'''
+'''FeatureFusionModule'''
 class FeatureFusionModule(nn.Module):
     def __init__(self, in_channels, out_channels, norm_cfg=None, act_cfg=None):
         super(FeatureFusionModule, self).__init__()

ssseg/modules/models/backbones/bisenetv2.py

@@ -17,7 +17,7 @@
 AUTO_ASSERT_STRUCTURE_TYPES = {}
 
 
-'''Detail Branch with wide channels and shallow layers to capture low-level details and generate high-resolution feature representation'''
+'''DetailBranch'''
 class DetailBranch(nn.Module):
     def __init__(self, detail_channels=(64, 64, 128), in_channels=3, norm_cfg=None, act_cfg=None):
         super(DetailBranch, self).__init__()
@@ -52,7 +52,7 @@ def forward(self, x):
         return x
 
 
-'''Stem Block at the beginning of Semantic Branch'''
+'''StemBlock'''
 class StemBlock(nn.Module):
     def __init__(self, in_channels=3, out_channels=16, norm_cfg=None, act_cfg=None):
         super(StemBlock, self).__init__()
@@ -84,7 +84,7 @@ def forward(self, x):
         return x
 
 
-'''Gather-and-Expansion Layer'''
+'''GELayer'''
 class GELayer(nn.Module):
     def __init__(self, in_channels, out_channels, exp_ratio=6, stride=1, norm_cfg=None, act_cfg=None):
         super(GELayer, self).__init__()
@@ -110,15 +110,8 @@ def __init__(self, in_channels, out_channels, exp_ratio=6, stride=1, norm_cfg=No
                 BuildActivation(act_cfg),
             )
             self.shortcut = nn.Sequential(DepthwiseSeparableConv2d(
-                in_channels=in_channels,
-                out_channels=out_channels,
-                kernel_size=3,
-                stride=stride,
-                padding=1,
-                dw_norm_cfg=norm_cfg,
-                dw_act_cfg=None,
-                pw_norm_cfg=norm_cfg,
-                pw_act_cfg=None,
+                in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=stride, padding=1,
+                dw_norm_cfg=norm_cfg, dw_act_cfg=None, pw_norm_cfg=norm_cfg, pw_act_cfg=None,
             ))
         self.conv2 = nn.Sequential(
             nn.Conv2d(mid_channel, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
@@ -140,7 +133,7 @@ def forward(self, x):
         return x
 
 
-'''Context Embedding Block for large receptive filed in Semantic Branch'''
+'''CEBlock'''
 class CEBlock(nn.Module):
     def __init__(self, in_channels=3, out_channels=16, norm_cfg=None, act_cfg=None):
         super(CEBlock, self).__init__()
@@ -172,7 +165,7 @@ def forward(self, x):
         return x
 
 
-'''Semantic Branch which is lightweight with narrow channels and deep layers to obtain high-level semantic context'''
+'''SemanticBranch'''
 class SemanticBranch(nn.Module):
     def __init__(self, semantic_channels=(16, 32, 64, 128), in_channels=3, exp_ratio=6, norm_cfg=None, act_cfg=None):
         super(SemanticBranch, self).__init__()
@@ -187,25 +180,18 @@ def __init__(self, semantic_channels=(16, 32, 64, 128), in_channels=3, exp_ratio
             if i == 0:
                 self.add_module(stage_name, StemBlock(in_channels, semantic_channels[i], norm_cfg=norm_cfg, act_cfg=act_cfg))
             elif i == (len(semantic_channels) - 1):
-                self.add_module(
-                    stage_name, 
-                    nn.Sequential(
-                        GELayer(semantic_channels[i - 1], semantic_channels[i], exp_ratio, 2, norm_cfg=norm_cfg, act_cfg=act_cfg),
-                        GELayer(semantic_channels[i], semantic_channels[i], exp_ratio, 1, norm_cfg=norm_cfg, act_cfg=act_cfg),
-                        GELayer(semantic_channels[i], semantic_channels[i], exp_ratio, 1, norm_cfg=norm_cfg, act_cfg=act_cfg),
-                        GELayer(semantic_channels[i], semantic_channels[i], exp_ratio, 1, norm_cfg=norm_cfg, act_cfg=act_cfg),
-                    )
-                )
+                self.add_module(stage_name, nn.Sequential(
+                    GELayer(semantic_channels[i - 1], semantic_channels[i], exp_ratio, 2, norm_cfg=norm_cfg, act_cfg=act_cfg),
+                    GELayer(semantic_channels[i], semantic_channels[i], exp_ratio, 1, norm_cfg=norm_cfg, act_cfg=act_cfg),
+                    GELayer(semantic_channels[i], semantic_channels[i], exp_ratio, 1, norm_cfg=norm_cfg, act_cfg=act_cfg),
+                    GELayer(semantic_channels[i], semantic_channels[i], exp_ratio, 1, norm_cfg=norm_cfg, act_cfg=act_cfg),
+                ))
             else:
-                self.add_module(
-                    stage_name,
-                    nn.Sequential(
-                        GELayer(semantic_channels[i - 1], semantic_channels[i], exp_ratio, 2, norm_cfg=norm_cfg, act_cfg=act_cfg),
-                        GELayer(semantic_channels[i], semantic_channels[i], exp_ratio, 1, norm_cfg=norm_cfg, act_cfg=act_cfg)
-                    )
-                )
-        self.add_module(
-            f'stage{len(semantic_channels)}_CEBlock',
+                self.add_module(stage_name, nn.Sequential(
+                    GELayer(semantic_channels[i - 1], semantic_channels[i], exp_ratio, 2, norm_cfg=norm_cfg, act_cfg=act_cfg),
+                    GELayer(semantic_channels[i], semantic_channels[i], exp_ratio, 1, norm_cfg=norm_cfg, act_cfg=act_cfg)
+                ))
+        self.add_module(f'stage{len(semantic_channels)}_CEBlock',
             CEBlock(semantic_channels[-1], semantic_channels[-1], norm_cfg=norm_cfg, act_cfg=act_cfg),
         )
         self.semantic_stages.append(f'stage{len(semantic_channels)}_CEBlock')
@@ -219,7 +205,7 @@ def forward(self, x):
         return semantic_outs
 
 
-'''Bilateral Guided Aggregation Layer to fuse the complementary information from both Detail Branch and Semantic Branch'''
+'''BGALayer'''
 class BGALayer(nn.Module):
     def __init__(self, out_channels=128, align_corners=False, norm_cfg=None, act_cfg=None):
         super(BGALayer, self).__init__()
@@ -228,15 +214,8 @@ def __init__(self, out_channels=128, align_corners=False, norm_cfg=None, act_cfg
         self.align_corners = align_corners
         # define modules
         self.detail_dwconv = nn.Sequential(DepthwiseSeparableConv2d(
-            in_channels=out_channels,
-            out_channels=out_channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            dw_norm_cfg=norm_cfg,
-            dw_act_cfg=None,
-            pw_norm_cfg=None,
-            pw_act_cfg=None,
+            in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1,
+            dw_norm_cfg=norm_cfg, dw_act_cfg=None, pw_norm_cfg=None, pw_act_cfg=None,
         ))
         self.detail_down = nn.Sequential(
             nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1, bias=False),
@@ -248,15 +227,8 @@ def __init__(self, out_channels=128, align_corners=False, norm_cfg=None, act_cfg
             BuildNormalization(placeholder=out_channels, norm_cfg=norm_cfg),
         )
         self.semantic_dwconv = nn.Sequential(DepthwiseSeparableConv2d(
-            in_channels=out_channels,
-            out_channels=out_channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-            dw_norm_cfg=norm_cfg,
-            dw_act_cfg=None,
-            pw_norm_cfg=None,
-            pw_act_cfg=None,
+            in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1,
+            dw_norm_cfg=norm_cfg, dw_act_cfg=None, pw_norm_cfg=None, pw_act_cfg=None,
         ))
         self.conv = nn.Sequential(
             nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
@@ -269,14 +241,10 @@ def forward(self, x_d, x_s):
         detail_down = self.detail_down(x_d)
         semantic_conv = self.semantic_conv(x_s)
         semantic_dwconv = self.semantic_dwconv(x_s)
-        semantic_conv = F.interpolate(
-            semantic_conv, size=detail_dwconv.shape[2:], mode='bilinear', align_corners=self.align_corners,
-        )
+        semantic_conv = F.interpolate(semantic_conv, size=detail_dwconv.shape[2:], mode='bilinear', align_corners=self.align_corners)
         fuse_1 = detail_dwconv * torch.sigmoid(semantic_conv)
         fuse_2 = detail_down * torch.sigmoid(semantic_dwconv)
-        fuse_2 = F.interpolate(
-            fuse_2, size=fuse_1.shape[2:], mode='bilinear', align_corners=self.align_corners
-        )
+        fuse_2 = F.interpolate(fuse_2, size=fuse_1.shape[2:], mode='bilinear', align_corners=self.align_corners)
         output = self.conv(fuse_1 + fuse_2)
         return output
 

ssseg/modules/models/backbones/cgnet.py

@@ -16,7 +16,7 @@
 AUTO_ASSERT_STRUCTURE_TYPES = {}
 
 
-'''Global Context Extractor for CGNet'''
+'''GlobalContextExtractor'''
 class GlobalContextExtractor(nn.Module):
     def __init__(self, channels, reduction=16):
         super(GlobalContextExtractor, self).__init__()
@@ -38,7 +38,7 @@ def forward(self, x):
         return x * y
 
 
-'''Context Guided Block for CGNet'''
+'''ContextGuidedBlock'''
 class ContextGuidedBlock(nn.Module):
     def __init__(self, in_channels, out_channels, dilation=2, reduction=16, skip_connect=True, downsample=False, norm_cfg=None, act_cfg=None):
         super(ContextGuidedBlock, self).__init__()
@@ -80,7 +80,7 @@ def forward(self, x):
         return out
 
 
-'''Downsampling module for CGNet'''
+'''InputInjection'''
 class InputInjection(nn.Module):
     def __init__(self, num_downsamplings):
         super(InputInjection, self).__init__()
@@ -144,14 +144,8 @@ def __init__(self, structure_type, in_channels=3, num_channels=(32, 64, 128), nu
         self.level1 = nn.ModuleList()
         for i in range(num_blocks[0]):
             self.level1.append(ContextGuidedBlock(
-                in_channels=cur_channels if i == 0 else num_channels[1], 
-                out_channels=num_channels[1], 
-                dilation=dilations[0], 
-                reduction=reductions[0], 
-                skip_connect=True, 
-                downsample=(i == 0), 
-                norm_cfg=norm_cfg, 
-                act_cfg=act_cfg,
+                in_channels=cur_channels if i == 0 else num_channels[1], out_channels=num_channels[1], dilation=dilations[0], 
+                reduction=reductions[0], skip_connect=True, downsample=(i == 0), norm_cfg=norm_cfg, act_cfg=act_cfg,
             ))
         cur_channels = 2 * num_channels[1] + in_channels
         self.norm_prelu_1 = nn.Sequential(
@@ -162,14 +156,8 @@ def __init__(self, structure_type, in_channels=3, num_channels=(32, 64, 128), nu
         self.level2 = nn.ModuleList()
         for i in range(num_blocks[1]):
             self.level2.append(ContextGuidedBlock(
-                in_channels=cur_channels if i == 0 else num_channels[2],
-                out_channels=num_channels[2],
-                dilation=dilations[1],
-                reduction=reductions[1],
-                skip_connect=True, 
-                downsample=(i == 0), 
-                norm_cfg=norm_cfg, 
-                act_cfg=act_cfg,
+                in_channels=cur_channels if i == 0 else num_channels[2], out_channels=num_channels[2], dilation=dilations[1],
+                reduction=reductions[1], skip_connect=True, downsample=(i == 0), norm_cfg=norm_cfg, act_cfg=act_cfg,
             ))
         cur_channels = 2 * num_channels[2]
         self.norm_prelu_2 = nn.Sequential(