18 lines
831 B
Python
18 lines
831 B
Python
|
from attention.ExternalAttention import ExternalAttention
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
input=torch.randn(50,49,512)
|
|||
|
|
ea = ExternalAttention(d_model=512,S=8)
|
|||
|
|
output=ea(input)
|
|||
|
|
print(output.shape)
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
1.1. 引用
|
|||
|
|
Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks.---arXiv 2021.05.05
|
|||
|
|
|
|||
|
|
论文地址：https://arxiv.org/abs/2105.02358
|
|||
|
|
这是五月份在arXiv上的一篇文章，主要解决的Self-Attention(SA)的两个痛点问题：
|
|||
|
|
（1）O(n^2)的计算复杂度；(2)SA是在同一个样本上根据不同位置计算Attention，忽略了不同样本之间的联系。
|
|||
|
|
因此，本文采用了两个串联的MLP结构作为memory units，使得计算复杂度降低到了O(n)；
|
|||
|
|
此外，这两个memory units是基于全部的训练数据学习的，因此也隐式的考虑了不同样本之间的联系。
|
|||
|
|
"""
|