Update networks/decoder.py

networks/decoder.py

@@ -1,48 +1,37 @@
-from typing import Dict, List, Optional, Sequence, Union
-
 import torch
 import torch.nn as nn
-from torch import Tensor
-import torch.nn.functional as F
 
-from mmcv.cnn import ConvModule
 from .height_head import resize
 
+
+class ConvModule(nn.Module):
+    """Conv + Norm + Activation with same submodule names as mmcv.ConvModule."""
+    def __init__(self, in_channels, out_channels, kernel_size, padding=0,
+                 norm_layer=None, act_layer=None):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,
+                              padding=padding, bias=(norm_layer is None))
+        self.bn = norm_layer(out_channels) if norm_layer is not None else nn.Identity()
+        self.activate = act_layer() if act_layer is not None else nn.Identity()
+
+    def forward(self, x):
+        return self.activate(self.bn(self.conv(x)))
+
 class PPM(nn.ModuleList):
-    """Pooling Pyramid Module used in PSPNet.
-    Args:
-        pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid
-            Module.
-        in_channels (int): Input channels.
-        channels (int): Channels after modules, before conv_seg.
-        conv_cfg (dict|None): Config of conv layers.
-        norm_cfg (dict|None): Config of norm layers.
-        act_cfg (dict): Config of activation layers.
-        align_corners (bool): align_corners argument of F.interpolate.
-    """
-
-    def __init__(self, pool_scales, in_channels, channels, conv_cfg, norm_cfg,
-                 act_cfg, align_corners, **kwargs):
+    def __init__(self, pool_scales, in_channels, channels, align_corners):
         super(PPM, self).__init__()
         self.pool_scales = pool_scales
         self.align_corners = align_corners
         self.in_channels = in_channels
         self.channels = channels
-        self.conv_cfg = conv_cfg
-        self.norm_cfg = norm_cfg
-        self.act_cfg = act_cfg
         for pool_scale in pool_scales:
             self.append(
                 nn.Sequential(
                     nn.AdaptiveAvgPool2d(pool_scale),
-                    ConvModule(
-                        self.in_channels,
-                        self.channels,
-                        1,
-                        conv_cfg=self.conv_cfg,
-                        norm_cfg=self.norm_cfg,
-                        act_cfg=self.act_cfg,
-                        **kwargs)))
+                    ConvModule(in_channels, channels, 1,
+                               norm_layer=nn.SyncBatchNorm,
+                               act_layer=nn.ReLU),
+                ))
 
     def forward(self, x):
         """Forward function."""
@@ -64,10 +53,8 @@ class Decoder(nn.Module):
                  num_deconv_filters=2,
                  decover_filter=[128, 128],
                  psp_channel=16,
-                 pool_scales=[1, 2, 3, 6],
-                 init_cfg=dict(
-                        type='TruncNormal', std=0.02, layer=['Conv2d', 'Linear']), ):
-        super(Decoder, self).__init__(init_cfg=init_cfg)
+                 pool_scales=[1, 2, 3, 6]):
+        super(Decoder, self).__init__()
 
         self.in_channel = in_channel
         self.num_deconv_filters = num_deconv_filters
@@ -114,9 +101,6 @@ class Decoder(nn.Module):
                 pool_scales=pool_scales,
                 in_channels=in_channel,
                 channels=psp_channel,
-                conv_cfg=dict(type='Conv2d'),
-                norm_cfg=dict(type='SyncBN', requires_grad=True),
-                act_cfg=dict(type='ReLU'),
                 align_corners=False)
 
     def _make_deconv_layer(self, name, in_channel, out_channel):