如何解决如何加快pytorch中的滑动注意力计算?
我正在尝试实现一个使用滑动注意窗口的注意模型。
到目前为止,我通过使用火炬模块 nn.MultiheadAttention 并传递一个注意力掩码来实现这一点,这允许每个位置 n 只关注来自 (n-W,n].
的数据
然而它很慢,我觉得在执行速度方面应该有改进的空间,因为滑动窗口只有实际系列大小的 1/10 左右。
鉴于 T 是我的序列长度和 W 滑动窗口,我觉得应该可以将执行时间从 O(T²) 减少到 O(W*T)。这个(玩具)功能已经做了我想要的,
def custom_matmul(a,b,mask):
device = a.device
B,T,_ = a.shape
ret = torch.zeros((B,T)).to(device)
for row,col in (mask==0).nonzero(as_tuple=False):
ret[:,row,col] = (a[:,row]*b[:,:,col]).sum(dim=-1)
return ret
但它仍然比在整个序列 T 上使用 a.matmul(b) 并在之后应用掩码要慢得多。有谁知道更快的实现方式?
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。