Some errors when using PSA

[ALmagical]

I wanted to insert PSA module to FCOS, but I met some errors.
Fortunately, I has solved that. I will show the changes I did.
There are three changes. First one is in the init function.

def __init__(self, channel=512,reduction=4,S=4):
       super().__init__()
       self.S=S

       self.convs=nn.ModuleList([])
       for i in range(S):
           # Add groups
           self.convs.append(nn.Conv2d(channel//S,channel//S,kernel_size=2*(i+1)+1,padding=i+1,groups=2**i))

       self.se_blocks=nn.ModuleList([])
       for i in range(S):
           self.se_blocks.append(nn.Sequential(
               nn.AdaptiveAvgPool2d(1),
               nn.Conv2d(channel//S, channel // (S*reduction),kernel_size=1, bias=False),
               nn.ReLU(inplace=False),
               nn.Conv2d(channel // (S*reduction), channel//S,kernel_size=1, bias=False),
               nn.Sigmoid()
           ))

       self.softmax=nn.Softmax(dim=1)

I find there are different groups setting for different SPC convs in the original EPSANet , so I add groups to same places. But I am not sure it is right.
Second, it is also in this function, I change self.convs=[] and self.se_blocks=[] to self.convs=nn.ModuleList([]) and self.se_blocks=nn.ModuleList([]). That is because when I used GPU to train my model, if their type are list, the initial weight will not be load to GPU. Using nn.ModuleList() will fix that.
Third, it is in the forward function. When I trained model, I got this error.
one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [1, 128, 32, 32]], which is output 0 of SliceBackward, is at version 3; expected version 2 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!
After try many different ways, I found the reason. This error is caused by the variables are changed in the forward function. There is my code.
` def forward(self, x):
b, c, h, w = x.size()

    #Step1:SPC module
    PSA_input=x.view(b,self.S,c//self.S,h,w) #bs,s,ci,h,w
    outs=[]
    for idx,conv in enumerate(self.convs):
        SPC_input = PSA_input[:,idx,:,:,:]
        #SPC_out[:,idx,:,:,:]=se(SPC_input)
        outs.append(conv(SPC_input))
    SPC_out = torch.stack([out for out in outs],dim=1)

    #Step2:SE weight
    SE_out=torch.zeros_like(SPC_out)
    outs=[]
    for idx,se in enumerate(self.se_blocks):
        SE_input = SPC_out[:,idx,:,:,:]
        #SE_out[:,idx,:,:,:]=se(SE_input)
        outs.append(se(SE_input))

    SE_out =  torch.stack([out for out in outs],dim=1)
    
    #Step3:Softmax
    softmax_out=self.softmax(SE_out)

    #Step4:SPA
    PSA_out=SPC_out*softmax_out
    PSA_out=PSA_out.view(b,-1,h,w)

    return PSA_out`

Now, my model with PSA can be trained, but don not know the result till now.
I think these changes don't change the struct of PSA, but I am not sure. I hope you can check if these changes are right.
Finally, thanks for your working.

[jialeli1]

Hi. How about the result with your work? 👀

[ALmagical]

I add PSA model between FPN and FCOS head. And I has tried two implied: 1. Sharing the PSA model between all FPN out levels. 2. Building different PSA model for each FPN out.
But, I find no any improvement or damage for the model training in my dataset.
I still don't know why.

[jialeli1]

The script was updated 3 hours ago. Maybe you can try it now.

External-Attention-pytorch/attention/PolarizedSelfAttention.py

Line 42 in b6a8fe2

spatial_wq=self.softmax_spatial(spatial_wq)