diff --git a/README.md b/README.md index 2af40a2..71d9211 100644 --- a/README.md +++ b/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) diff --git a/attention/EMSA.py b/attention/EMSA.py new file mode 100644 index 0000000..b13dcec --- /dev/null +++ b/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) + + \ No newline at end of file diff --git a/attention/__pycache__/EMSA.cpython-38.pyc b/attention/__pycache__/EMSA.cpython-38.pyc new file mode 100644 index 0000000..9d39732 Binary files /dev/null and b/attention/__pycache__/EMSA.cpython-38.pyc differ diff --git a/img/EMSA.png b/img/EMSA.png new file mode 100644 index 0000000..32393b1 Binary files /dev/null and b/img/EMSA.png differ diff --git a/main.py b/main.py index 68700fe..74ad961 100644 --- a/main.py +++ b/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) \ No newline at end of file