PyTorch实现 Self Attention

文章预览

点击上方 “ 小白学视觉 ”，选择加" 星标 "或“ 置顶 ” 重磅干货，第一时间送达 仅作学术分享，不代表本公众号立场，侵权联系删除 转载于： 作者丨Connolly@知乎（已授权） 来源丨https://zhuanlan.zhihu.com/p/445016136 编辑丨极市平台 通过修改SelfAttention的执行逻辑，可以节省大量的激活值显存开销。 这篇文章的消除方法来自于2021年12月10日谷歌放到arxiv上的文章self attention does not need O(n^2) memory. 该方法巧妙地使用了小学学到的加法分配率，将self attention中的固定激活值降到了O(1)的程度。[1] Self Attention 固定激活值显存分析 Hugging face Transformers中，SelfAttention 内核实现 表格中只列举了会实测中产生激活值的操作，其中B为Batch_size，L为sequence_length，H为hidden_size，m为SelfAttention中head的数量。 则总和  。 观察： 当   固定时, 即模型结构是固定的时候, 我们发现 ………………………………