Spaces:
Runtime error
Runtime error
import numpy as np | |
import torch | |
from torch import nn | |
from torch.nn import init | |
from collections import OrderedDict | |
class SKAttention(nn.Module): | |
def __init__(self, channel=512, kernels=[1,3,5,7], reduction=16, group=1, L=32): | |
super().__init__() | |
self.d = max(L, channel // reduction) | |
self.convs = nn.ModuleList([]) | |
for k in kernels: | |
self.convs.append( | |
nn.Sequential(OrderedDict([ | |
('conv', nn.Conv2d(channel, channel, kernel_size=k, padding=k // 2, groups=group)), | |
('bn', nn.BatchNorm2d(channel)), | |
('relu', nn.ReLU()) | |
])) | |
) | |
self.fc = nn.Linear(channel, self.d) | |
self.fcs = nn.ModuleList([]) | |
for i in range(len(kernels)): | |
self.fcs.append(nn.Linear(self.d, channel)) | |
self.softmax = nn.Softmax(dim=0) | |
def forward(self, x): | |
bs, c, _, _ = x.size() | |
conv_outs = [] | |
### split | |
for conv in self.convs: | |
conv_outs.append(conv(x)) | |
feats = torch.stack(conv_outs, 0) # k,bs,channel,h,w | |
### fuse | |
U = sum(conv_outs) # bs,c,h,w | |
### reduction channel | |
S = U.mean(-1).mean(-1) # bs,c | |
Z = self.fc(S) # bs,d | |
### calculate attention weight | |
weights = [] | |
for fc in self.fcs: | |
weight = fc(Z) | |
weights.append(weight.view(bs, c, 1, 1)) # bs,channel | |
attention_weughts = torch.stack(weights, 0) # k,bs,channel,1,1 | |
attention_weughts = self.softmax(attention_weughts) # k,bs,channel,1,1 | |
### fuse | |
V = (attention_weughts * feats).sum(0) | |
return V |