update

regainOWO · Aug 17, 2022 · 88bd677 · 88bd677
1 parent 306f987
commit 88bd677
Show file tree

Hide file tree

Showing 3 changed files with 205 additions and 0 deletions.
diff --git a/configs/attention/yolov5_SOCA.yaml b/configs/attention/yolov5_SOCA.yaml
@@ -0,0 +1,49 @@
+# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
+
+# Parameters
+nc: 80  # number of classes
+depth_multiple: 0.33  # model depth multiple
+width_multiple: 0.50  # layer channel multiple
+anchors:
+  - [10,13, 16,30, 33,23]  # P3/8
+  - [30,61, 62,45, 59,119]  # P4/16
+  - [116,90, 156,198, 373,326]  # P5/32
+
+# YOLOv5 v6.0 backbone
+backbone:
+  # [from, number, module, args]
+  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
+   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
+   [-1, 3, C3, [128]],
+   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
+   [-1, 6, C3, [256]],
+   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
+   [-1, 9, C3, [512]],
+   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
+   [-1, 3, C3, [1024]],
+   [-1, 1, SPPF, [1024, 5]],  # 9
+  ]
+
+# YOLOv5 v6.0 head
+head:
+  [[-1, 1, Conv, [512, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
+   [-1, 3, C3, [512, False]],  # 13
+
+   [-1, 1, Conv, [256, 1, 1]],
+   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
+   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
+   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
+
+   [-1, 1, Conv, [256, 3, 2]],
+   [[-1, 14], 1, Concat, [1]],  # cat head P4
+   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
+
+   [-1, 1, Conv, [512, 3, 2]],
+   [[-1, 10], 1, Concat, [1]],  # cat head P5
+   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
+   [-1, 1, SOCA, [512]],
+
+   [[17, 20, 24], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  ]
diff --git a/models/Models/Attention/SOCA.py b/models/Models/Attention/SOCA.py
@@ -0,0 +1,154 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import init
+
+from torch.autograd import Function
+
+class Covpool(Function):
+     @staticmethod
+     def forward(ctx, input):
+         x = input
+         batchSize = x.data.shape[0]
+         dim = x.data.shape[1]
+         h = x.data.shape[2]
+         w = x.data.shape[3]
+         M = h*w
+         x = x.reshape(batchSize,dim,M)
+         I_hat = (-1./M/M)*torch.ones(M,M,device = x.device) + (1./M)*torch.eye(M,M,device = x.device)
+         I_hat = I_hat.view(1,M,M).repeat(batchSize,1,1).type(x.dtype)
+         y = x.bmm(I_hat).bmm(x.transpose(1,2))
+         ctx.save_for_backward(input,I_hat)
+         return y
+     @staticmethod
+     def backward(ctx, grad_output):
+         input,I_hat = ctx.saved_tensors
+         x = input
+         batchSize = x.data.shape[0]
+         dim = x.data.shape[1]
+         h = x.data.shape[2]
+         w = x.data.shape[3]
+         M = h*w
+         x = x.reshape(batchSize,dim,M)
+         grad_input = grad_output + grad_output.transpose(1,2)
+         grad_input = grad_input.bmm(x).bmm(I_hat)
+         grad_input = grad_input.reshape(batchSize,dim,h,w)
+         return grad_input
+
+class Sqrtm(Function):
+     @staticmethod
+     def forward(ctx, input, iterN):
+         x = input
+         batchSize = x.data.shape[0]
+         dim = x.data.shape[1]
+         dtype = x.dtype
+         I3 = 3.0*torch.eye(dim,dim,device = x.device).view(1, dim, dim).repeat(batchSize,1,1).type(dtype)
+         normA = (1.0/3.0)*x.mul(I3).sum(dim=1).sum(dim=1)
+         A = x.div(normA.view(batchSize,1,1).expand_as(x))
+         Y = torch.zeros(batchSize, iterN, dim, dim, requires_grad = False, device = x.device)
+         Z = torch.eye(dim,dim,device = x.device).view(1,dim,dim).repeat(batchSize,iterN,1,1)
+         if iterN < 2:
+            ZY = 0.5*(I3 - A)
+            Y[:,0,:,:] = A.bmm(ZY)
+         else:
+            ZY = 0.5*(I3 - A)
+            Y[:,0,:,:] = A.bmm(ZY)
+            Z[:,0,:,:] = ZY
+            for i in range(1, iterN-1):
+               ZY = 0.5*(I3 - Z[:,i-1,:,:].bmm(Y[:,i-1,:,:]))
+               Y[:,i,:,:] = Y[:,i-1,:,:].bmm(ZY)
+               Z[:,i,:,:] = ZY.bmm(Z[:,i-1,:,:])
+            ZY = 0.5*Y[:,iterN-2,:,:].bmm(I3 - Z[:,iterN-2,:,:].bmm(Y[:,iterN-2,:,:]))
+         y = ZY*torch.sqrt(normA).view(batchSize, 1, 1).expand_as(x)
+         ctx.save_for_backward(input, A, ZY, normA, Y, Z)
+         ctx.iterN = iterN
+         return y
+     @staticmethod
+     def backward(ctx, grad_output):
+         input, A, ZY, normA, Y, Z = ctx.saved_tensors
+         iterN = ctx.iterN
+         x = input
+         batchSize = x.data.shape[0]
+         dim = x.data.shape[1]
+         dtype = x.dtype
+         der_postCom = grad_output*torch.sqrt(normA).view(batchSize, 1, 1).expand_as(x)
+         der_postComAux = (grad_output*ZY).sum(dim=1).sum(dim=1).div(2*torch.sqrt(normA))
+         I3 = 3.0*torch.eye(dim,dim,device = x.device).view(1, dim, dim).repeat(batchSize,1,1).type(dtype)
+         if iterN < 2:
+            der_NSiter = 0.5*(der_postCom.bmm(I3 - A) - A.bmm(der_sacleTrace))
+         else:
+            dldY = 0.5*(der_postCom.bmm(I3 - Y[:,iterN-2,:,:].bmm(Z[:,iterN-2,:,:])) -
+                          Z[:,iterN-2,:,:].bmm(Y[:,iterN-2,:,:]).bmm(der_postCom))
+            dldZ = -0.5*Y[:,iterN-2,:,:].bmm(der_postCom).bmm(Y[:,iterN-2,:,:])
+            for i in range(iterN-3, -1, -1):
+               YZ = I3 - Y[:,i,:,:].bmm(Z[:,i,:,:])
+               ZY = Z[:,i,:,:].bmm(Y[:,i,:,:])
+               dldY_ = 0.5*(dldY.bmm(YZ) - 
+                         Z[:,i,:,:].bmm(dldZ).bmm(Z[:,i,:,:]) - 
+                             ZY.bmm(dldY))
+               dldZ_ = 0.5*(YZ.bmm(dldZ) - 
+                         Y[:,i,:,:].bmm(dldY).bmm(Y[:,i,:,:]) -
+                            dldZ.bmm(ZY))
+               dldY = dldY_
+               dldZ = dldZ_
+            der_NSiter = 0.5*(dldY.bmm(I3 - A) - dldZ - A.bmm(dldY))
+         grad_input = der_NSiter.div(normA.view(batchSize,1,1).expand_as(x))
+         grad_aux = der_NSiter.mul(x).sum(dim=1).sum(dim=1)
+         for i in range(batchSize):
+             grad_input[i,:,:] += (der_postComAux[i] \
+                                   - grad_aux[i] / (normA[i] * normA[i])) \
+                                   *torch.ones(dim,device = x.device).diag()
+         return grad_input, None
+
+def CovpoolLayer(var):
+    return Covpool.apply(var)
+
+def SqrtmLayer(var, iterN):
+    return Sqrtm.apply(var, iterN)
+
+class SOCA(nn.Module):
+    # second-order Channel attention
+    def __init__(self, channel, reduction=8):
+        super(SOCA, self).__init__()
+        self.max_pool = nn.MaxPool2d(kernel_size=2)
+
+        self.conv_du = nn.Sequential(
+            nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        batch_size, C, h, w = x.shape  # x: NxCxHxW
+        N = int(h * w)
+        min_h = min(h, w)
+        h1 = 1000
+        w1 = 1000
+        if h < h1 and w < w1:
+            x_sub = x
+        elif h < h1 and w > w1:
+            W = (w - w1) // 2
+            x_sub = x[:, :, :, W:(W + w1)]
+        elif w < w1 and h > h1:
+            H = (h - h1) // 2
+            x_sub = x[:, :, H:H + h1, :]
+        else:
+            H = (h - h1) // 2
+            W = (w - w1) // 2
+            x_sub = x[:, :, H:(H + h1), W:(W + w1)]
+        cov_mat = CovpoolLayer(x_sub) # Global Covariance pooling layer
+        cov_mat_sqrt = SqrtmLayer(cov_mat,5) # Matrix square root layer( including pre-norm,Newton-Schulz iter. and post-com. with 5 iteration)
+        cov_mat_sum = torch.mean(cov_mat_sqrt,1)
+        cov_mat_sum = cov_mat_sum.view(batch_size,C,1,1)
+        y_cov = self.conv_du(cov_mat_sum)
+        return y_cov*x
+
+if __name__ == '__main__':
+    # https://github.com/daitao/SAN
+    input=torch.randn(50,512,7,7)
+    kernel_size=input.shape[2]
+    cbam = SOCA(channel=512)
+    output=cbam(input)
+    print(output.shape)
+
diff --git a/models/common.py b/models/common.py
@@ -47,6 +47,8 @@
 from models.Models.Attention.S2Attention import S2Attention
 from models.Models.Attention.SEAttention import SEAttention
 from models.Models.Attention.SKAttention import SKAttention
+from models.Models.Attention.SOCA import SOCA
+