mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
f902bcd54c
* select_conv2d -> create_conv2d * added create_attn to create attention module from string/bool/module * factor padding helpers into own file, use in both conv2d_same and avg_pool2d_same * add some more test eca resnet variants * minor tweaks, naming, comments, consistency
22 lines
716 B
Python
22 lines
716 B
Python
from torch import nn as nn
|
|
|
|
|
|
class SEModule(nn.Module):
|
|
|
|
def __init__(self, channels, reduction=16, act_layer=nn.ReLU):
|
|
super(SEModule, self).__init__()
|
|
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
|
reduction_channels = max(channels // reduction, 8)
|
|
self.fc1 = nn.Conv2d(
|
|
channels, reduction_channels, kernel_size=1, padding=0, bias=True)
|
|
self.act = act_layer(inplace=True)
|
|
self.fc2 = nn.Conv2d(
|
|
reduction_channels, channels, kernel_size=1, padding=0, bias=True)
|
|
|
|
def forward(self, x):
|
|
x_se = self.avg_pool(x)
|
|
x_se = self.fc1(x_se)
|
|
x_se = self.act(x_se)
|
|
x_se = self.fc2(x_se)
|
|
return x * x_se.sigmoid()
|