|
|
|
import mmcv |
|
import torch.nn as nn |
|
from mmcv.cnn import ConvModule |
|
|
|
|
|
class SELayer(nn.Module): |
|
"""Squeeze-and-Excitation Module. |
|
|
|
Args: |
|
channels (int): The input (and output) channels of the SE layer. |
|
ratio (int): Squeeze ratio in SELayer, the intermediate channel will be |
|
``int(channels/ratio)``. Default: 16. |
|
conv_cfg (None or dict): Config dict for convolution layer. |
|
Default: None, which means using conv2d. |
|
act_cfg (dict or Sequence[dict]): Config dict for activation layer. |
|
If act_cfg is a dict, two activation layers will be configurated |
|
by this dict. If act_cfg is a sequence of dicts, the first |
|
activation layer will be configurated by the first dict and the |
|
second activation layer will be configurated by the second dict. |
|
Default: (dict(type='ReLU'), dict(type='Sigmoid')) |
|
""" |
|
|
|
def __init__(self, |
|
channels, |
|
ratio=16, |
|
conv_cfg=None, |
|
act_cfg=(dict(type='ReLU'), dict(type='Sigmoid'))): |
|
super().__init__() |
|
if isinstance(act_cfg, dict): |
|
act_cfg = (act_cfg, act_cfg) |
|
assert len(act_cfg) == 2 |
|
assert mmcv.is_tuple_of(act_cfg, dict) |
|
self.global_avgpool = nn.AdaptiveAvgPool2d(1) |
|
self.conv1 = ConvModule( |
|
in_channels=channels, |
|
out_channels=int(channels / ratio), |
|
kernel_size=1, |
|
stride=1, |
|
conv_cfg=conv_cfg, |
|
act_cfg=act_cfg[0]) |
|
self.conv2 = ConvModule( |
|
in_channels=int(channels / ratio), |
|
out_channels=channels, |
|
kernel_size=1, |
|
stride=1, |
|
conv_cfg=conv_cfg, |
|
act_cfg=act_cfg[1]) |
|
|
|
def forward(self, x): |
|
out = self.global_avgpool(x) |
|
out = self.conv1(out) |
|
out = self.conv2(out) |
|
return x * out |
|
|
