# Codeblock 3
class SEModule(nn.Module):
    def __init__(self, num_channels, r):
        super().__init__()
        
        self.global_pooling = nn.AdaptiveAvgPool2d(output_size=(1,1))
        self.fc0 = nn.Linear(in_features=num_channels,
                             out_features=num_channels//r, 
                             bias=False)
        self.relu6 = nn.ReLU6()
        self.fc1 = nn.Linear(in_features=num_channels//r,
                             out_features=num_channels, 
                             bias=False)
        self.hardsigmoid = nn.Hardsigmoid()

    def forward(self, x):
        print(f'original\t\t: {x.size()}')
        
        squeezed = self.global_pooling(x)              #(1)
        print(f'after avgpool\t\t: {squeezed.size()}')
        
        squeezed = torch.flatten(squeezed, 1)
        print(f'after flatten\t\t: {squeezed.size()}')
        
        excited = self.fc0(squeezed)                   #(2)
        print(f'after fc0\t\t: {excited.size()}')
        
        excited = self.relu6(excited)
        print(f'after relu6\t\t: {excited.size()}')
        
        excited = self.fc1(excited)                    #(3)
        print(f'after fc1\t\t: {excited.size()}')
        
        excited = self.hardsigmoid(excited)            #(4)
        print(f'after hardsigmoid\t: {excited.size()}')
        
        excited = excited[:, :, None, None]
        print(f'after reshape\t\t: {excited.size()}')
        
        scaled = x * excited                           #(5)
        print(f'after scaling\t\t: {scaled.size()}')
        
        return scaled