Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class h_sigmoid(nn.Module): | |
| def __init__(self, inplace=True): | |
| super(h_sigmoid, self).__init__() | |
| self.relu = nn.ReLU6(inplace=inplace) | |
| def forward(self, x): | |
| return self.relu(x + 3) / 6 | |
| class h_swish(nn.Module): | |
| def __init__(self, inplace=True): | |
| super(h_swish, self).__init__() | |
| self.sigmoid = h_sigmoid(inplace=inplace) | |
| def forward(self, x): | |
| return x * self.sigmoid(x) | |
| class CoordAtt(nn.Module): | |
| def __init__(self, inp, oup, reduction=32): | |
| super(CoordAtt, self).__init__() | |
| self.pool_h = nn.AdaptiveAvgPool2d((None, 1)) | |
| self.pool_w = nn.AdaptiveAvgPool2d((1, None)) | |
| mip = max(8, inp // reduction) | |
| self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0) | |
| self.bn1 = nn.BatchNorm2d(mip) | |
| self.bn2 = nn.BatchNorm2d(1) | |
| self.bn3 = nn.BatchNorm2d(1) | |
| self.act = h_swish() | |
| self.bn4 = nn.BatchNorm2d(mip) | |
| self.bn5 = nn.BatchNorm2d(mip) | |
| self.bn6 = nn.BatchNorm2d(1) | |
| self.bn7 = nn.BatchNorm2d(1) | |
| self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0) | |
| self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0) | |
| def forward(self, x): | |
| x = torch.unsqueeze(x, 1) #2 1 2304 196 | |
| identity = x | |
| n, c, h, w = x.size()#2 1 2304 196 | |
| x_h = self.bn2(self.pool_h(x))#2 1 2304 1 | |
| x_w = self.bn3(self.pool_w(x).permute(0, 1, 3, 2)) #2 1 196 1 | |
| identity_x_w = x_w | |
| identity_x_h = x_h | |
| y = torch.cat([x_h, x_w], dim=2) | |
| y = self.conv1(y) #2 8 2500 1 | |
| y = self.bn1(y) | |
| y = self.act(y) | |
| x_h, x_w = torch.split(y, [h, w], dim=2) #2 8 2304 1 | 2 8 196 1 | |
| x_h = self.bn4(x_h)+identity_x_h | |
| x_w = self.bn5(x_w)+identity_x_w | |
| x_w = x_w.permute(0, 1, 3, 2) | |
| a_h = self.bn6(self.conv_h(x_h)).sigmoid() #2 1 2304 1 | |
| a_w = self.bn7(self.conv_w(x_w)).sigmoid() #24 1 1 196 | |
| out = identity * a_w * a_h #点× | |
| out = torch.squeeze(out, 1) | |
| return out | |
| class CoordAtt_ori(nn.Module): | |
| def __init__(self, inp, oup, reduction=32): | |
| super(CoordAtt_ori, self).__init__() | |
| self.pool_h = nn.AdaptiveAvgPool2d((None, 1)) | |
| self.pool_w = nn.AdaptiveAvgPool2d((1, None)) | |
| mip = max(8, inp // reduction) | |
| self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0) | |
| self.bn1 = nn.BatchNorm2d(mip) | |
| self.act = h_swish() | |
| self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0) | |
| self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0) | |
| def forward(self, x): | |
| x = torch.unsqueeze(x, 1) | |
| identity = x | |
| n, c, h, w = x.size() | |
| x_h = self.pool_h(x) | |
| x_w = self.pool_w(x).permute(0, 1, 3, 2) | |
| y = torch.cat([x_h, x_w], dim=2) | |
| y = self.conv1(y) | |
| y = self.bn1(y) | |
| y = self.act(y) | |
| x_h, x_w = torch.split(y, [h, w], dim=2) | |
| x_w = x_w.permute(0, 1, 3, 2) | |
| a_h = self.conv_h(x_h).sigmoid() | |
| a_w = self.conv_w(x_w).sigmoid() | |
| out = identity * a_w * a_h | |
| out = torch.squeeze(out, 1) | |
| return out |