import torch.nn as nn
class SENet(nn.Module):
'''
input:(batch_size, chs, h, w)
output:(batch_size, chs, h, w)
'''
def __init__(self, chs, reduction=4):
super(SENet, self).__init__()
self.average_pooling = nn.AdaptiveAvgPool2d(output_size=(1, 1))
# (batch_size, chs, h, w) -> (batch, chs, 1, 1)
self.fc = nn.Sequential(
# First reduce dimension, then raise dimension.
# Add nonlinear processing to fit the correlation between channels
nn.Linear(chs, chs // reduction),
nn.ReLU(inplace=True),
nn.Linear(chs // reduction, chs)
)
self.activation = nn.Sigmoid() # Get the weight of each channel
def forward(self, x):
tmp = x
batch_size, chs, h, w = x.shape
x = self.average_pooling(x).view(batch_size, chs)
# (batch_size, chs, h, w) -> (batch, chs, 1, 1) -> (batch, chs)
x = self.fc(x).view(batch_size, chs, 1, 1)
# (batch, chs) -> (batch, chs, 1, 1)
# x = torch.clamp(x, min=0, max=1)
# if min =< elem <= max, new_elem = elem, if elem > max, new_elem = max, if elem < min, new_elem = min
x = self.activation(x)
return x * tmp # x: (batch_size, chs, 1, 1), tmp: (batch_size, chs, h, w)
