import numpy as np
import torch
from torch import nn
from torch.nn import init
class PSA(nn.Module):
def __init__(self, channel=512,reduction=4,S=4):
super().__init__()
self.S=S
self.convs=[]
for i in range(S):
self.convs.append(nn.Conv2d(channel//S,channel//S,kernel_size=2*(i+1)+1,padding=i+1))
self.se_blocks=[]
for i in range(S):#SEAttention
self.se_blocks.append(nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(channel//S, channel // (S*reduction),kernel_size=1, bias=False),#1*1 conv equals Fully connected layer
nn.ReLU(inplace=True),
nn.Conv2d(channel // (S*reduction), channel//S,kernel_size=1, bias=False),
nn.Sigmoid()
))
self.softmax=nn.Softmax(dim=1)
def init_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
init.normal_(m.weight, std=0.001)
if m.bias is not None:
init.constant_(m.bias, 0)
def forward(self, x):
b, c, h, w = x.size()
#Step1:SPC module
SPC_out=x.view(b,self.S,c//self.S,h,w) #bs,s,ci,h,w
for idx,conv in enumerate(self.convs):
print(SPC_out[:,idx,:,:,:].shape)
SPC_out[:,idx,:,:,:]=conv(SPC_out[:,idx,:,:,:])
print("SPC_out.shape:{:}".format(SPC_out[:,idx,:,:,:].shape))
#Step2:SE weight
se_out=[]
for idx,se in enumerate(self.se_blocks):
se_out.append(se(SPC_out[:,idx,:,:,:]))
SE_out=torch.stack(se_out,dim=1)
SE_out=SE_out.expand_as(SPC_out)
#Step3:Softmax
softmax_out=self.softmax(SE_out)
print("softmax_out.shape:{:}".format(softmax_out.shape))
#Step4:SPA
PSA_out=SPC_out*softmax_out
PSA_out=PSA_out.view(b,-1,h,w)
return PSA_out
if __name__ == '__main__':
input=torch.randn(50,512,7,7)
psa = PSA(channel=512,reduction=8)
output=psa(input)
print("output.shape:{:}".format(output.shape))
a=output.view(-1).sum()
print(a)
a.backward()
print(output.shape)
Pyramid Squeeze Attention
最新推荐文章于 2024-04-26 10:05:19 发布