modules: Refactor code

models/modules/aspp.py

@@ -8,9 +8,10 @@
 class ASPPConv(nn.Sequential):
  def __init__(self, in_channels, out_channels, dilation):
  modules = [
- models.modules.conv.SeparableConv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation,
- activation=nn.ReLU(),
- channel_attention=models.modules.attention.ChannelAttention(in_channels)),
+ models.modules.conv.SeparableConv2d(
+ in_channels, out_channels, 3, padding=dilation, dilation=dilation,
+ channel_attention=models.modules.attention.ChannelAttention(in_channels)
+ ),
  nn.ReLU()
  ]
  super(ASPPConv, self).__init__(*modules)
@@ -53,7 +54,6 @@ def __init__(self, in_channels, atrous_rates, out_channels=256):
  nn.BatchNorm2d(out_channels),
  nn.ReLU()
  )
-
  self.shortcut = nn.Sequential(
  nn.Conv2d(in_channels, out_channels, 1, bias=False),
  nn.BatchNorm2d(out_channels),

models/modules/attention.py

@@ -4,14 +4,11 @@
 
 class ChannelAttention(nn.Module):
  """ Squeeze and Excitation"""
- def __init__(self, in_channels: int, reduction_ratio=4, activation: nn.Module = None, multiplication=True):
+ def __init__(self, in_channels: int, reduction_ratio=4, activation=nn.ReLU(), multiplication=True):
  super(ChannelAttention, self).__init__()
  self.gap = nn.AdaptiveAvgPool2d(1)
  self.conv1 = nn.Conv2d(in_channels, in_channels // reduction_ratio, 1)
- if activation is not None:
- self.activation = activation
- else:
- self.activation = nn.ReLU()
+ self.activation = activation
  self.conv2 = nn.Conv2d(in_channels // reduction_ratio, in_channels, 1)
  self.sigmoid = nn.Sigmoid()
  self.multiplication = multiplication