EMSA

README.md

@@ -19,8 +19,9 @@
 
     - [9. DANet Attention Usage](#9-danet-attention-usage)
 
-    - [10. Pyramid Split Attention (PSA) Attention Usage](#10-Pyramid-Split-Attention-(PSA)-attention-usage)
+    - [10. Pyramid Split Attention (PSA) Attention Usage](#10-Pyramid-Split-Attention-(PSA)-Attention-Usage)
 
+    - [11. Efficient Multi-Head Self-Attention(EMSA) Usage](#11-Efficient-Multi-Head-Self-Attention(EMSA)-Usage)
 
 - [MLP Series](#mlp-series)
 
@@ -65,6 +66,10 @@
 
 - Pytorch implementation of ["EPSANet: An Efficient Pyramid Split Attention Block on Convolutional Neural Network---arXiv 2020.05.30"](https://arxiv.org/pdf/2105.14447.pdf)
 
+- Pytorch implementation of ["ResT: An Efficient Transformer for Visual Recognition---arXiv 2020.05.28"](https://arxiv.org/abs/2105.13677)
+
+
+
 ***
 
 
@@ -277,12 +282,37 @@ if __name__ == '__main__':
     output=psa(input)
     print(output.shape)
 
-
 ```
 
 ***
 
 
+### 11. Efficient Multi-Head Self-Attention(EMSA) Usage
+
+#### 11.1. Paper
+["ResT: An Efficient Transformer for Visual Recognition"](https://arxiv.org/abs/2105.13677)
+
+#### 11.2. Overview
+![](./img/EMSA.png)
+
+#### 11.3. Code
+```python
+
+from attention.EMSA import EMSA
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+if __name__ == '__main__':
+    input=torch.randn(50,64,512)
+    emsa = EMSA(d_model=512, d_k=512, d_v=512, h=8,H=8,W=8,ratio=2,apply_transform=True)
+    output=emsa(input,input,input)
+    print(output.shape)
+
+```
+
+***
+
 # MLP Series
 
 - Pytorch implementation of ["RepMLP: Re-parameterizing Convolutions into Fully-connected Layers for Image Recognition---arXiv 2020.05.05"](https://arxiv.org/pdf/2105.01883v1.pdf)

attention/EMSA.py

@@ -0,0 +1,100 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import init
+
+
+
+class EMSA(nn.Module):
+
+    def __init__(self, d_model, d_k, d_v, h,dropout=.1,H=7,W=7,ratio=3,apply_transform=True):
+
+        super(EMSA, self).__init__()
+        self.H=H
+        self.W=W
+        self.fc_q = nn.Linear(d_model, h * d_k)
+        self.fc_k = nn.Linear(d_model, h * d_k)
+        self.fc_v = nn.Linear(d_model, h * d_v)
+        self.fc_o = nn.Linear(h * d_v, d_model)
+        self.dropout=nn.Dropout(dropout)
+
+        self.ratio=ratio
+        if(self.ratio>1):
+            self.sr=nn.Sequential()
+            self.sr_conv=nn.Conv2d(d_model,d_model,kernel_size=ratio+1,stride=ratio,padding=ratio//2,groups=d_model)
+            self.sr_ln=nn.LayerNorm(d_model)
+
+        self.apply_transform=apply_transform and h>1
+        if(self.apply_transform):
+            self.transform=nn.Sequential()
+            self.transform.add_module('conv',nn.Conv2d(h,h,kernel_size=1,stride=1))
+            self.transform.add_module('softmax',nn.Softmax(-1))
+            self.transform.add_module('in',nn.InstanceNorm2d(h))
+
+        self.d_model = d_model
+        self.d_k = d_k
+        self.d_v = d_v
+        self.h = h
+
+        self.init_weights()
+
+
+    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, queries, keys, values, attention_mask=None, attention_weights=None):
+
+        b_s, nq ,c = queries.shape
+        nk = keys.shape[1]
+
+        q = self.fc_q(queries).view(b_s, nq, self.h, self.d_k).permute(0, 2, 1, 3)  # (b_s, h, nq, d_k)
+
+        if(self.ratio>1):
+            x=queries.permute(0,2,1).view(b_s,c,self.H,self.W) #bs,c,H,W
+            x=self.sr_conv(x) #bs,c,h,w
+            x=x.contiguous().view(b_s,c,-1).permute(0,2,1) #bs,n',c
+            x=self.sr_ln(x)
+            k = self.fc_k(x).view(b_s, -1, self.h, self.d_k).permute(0, 2, 3, 1)  # (b_s, h, d_k, n')
+            v = self.fc_v(x).view(b_s, -1, self.h, self.d_v).permute(0, 2, 1, 3)  # (b_s, h, n', d_v)
+        else:
+            k = self.fc_k(keys).view(b_s, nk, self.h, self.d_k).permute(0, 2, 3, 1)  # (b_s, h, d_k, nk)
+            v = self.fc_v(values).view(b_s, nk, self.h, self.d_v).permute(0, 2, 1, 3)  # (b_s, h, nk, d_v)
+
+        if(self.apply_transform):
+            att = torch.matmul(q, k) / np.sqrt(self.d_k)  # (b_s, h, nq, n')
+            att = self.transform(att) # (b_s, h, nq, n')
+        else:
+            att = torch.matmul(q, k) / np.sqrt(self.d_k)  # (b_s, h, nq, n')
+            att = torch.softmax(att, -1) # (b_s, h, nq, n')
+
+
+        if attention_weights is not None:
+            att = att * attention_weights
+        if attention_mask is not None:
+            att = att.masked_fill(attention_mask, -np.inf)
+
+        att=self.dropout(att)
+
+        out = torch.matmul(att, v).permute(0, 2, 1, 3).contiguous().view(b_s, nq, self.h * self.d_v)  # (b_s, nq, h*d_v)
+        out = self.fc_o(out)  # (b_s, nq, d_model)
+        return out
+
+
+if __name__ == '__main__':
+    input=torch.randn(50,64,512)
+    emsa = EMSA(d_model=512, d_k=512, d_v=512, h=8,H=8,W=8,ratio=2,apply_transform=True)
+    output=emsa(input,input,input)
+    print(output.shape)
+
+

main.py

@@ -1,16 +1,12 @@
-from rep.acnet import ACNet
+from attention.EMSA import EMSA
 import torch
 from torch import nn
 from torch.nn import functional as F
 
 if __name__ == '__main__':
-    input=torch.randn(50,512,49,49)
-    acnet=ACNet(512,512)
-    acnet.eval()
-    out=acnet(input)
-    acnet._switch_to_deploy()
-    out2=acnet(input)
-    print('difference:')
-    print(((out2-out)**2).sum())
+    input=torch.randn(50,64,512)
+    emsa = EMSA(d_model=512, d_k=512, d_v=512, h=8,H=8,W=8,ratio=2,apply_transform=True)
+    output=emsa(input,input,input)
+    print(output.shape)