from math import sqrt
import torch
import torch.nn as nnclass SA(nn.Module):def __init__(self, dimQ, dimK, dimV):super(SA, self).__init__()self.dimQ = dimQself.dimK = dimKself.dimV = dimV# self.mid = 10self.linerQ = nn.Linear(self.dimQ, self.dimV, bias = False)self.linerK = nn.Linear(self.dimK, self.dimV, bias = False)self.linerV = nn.Linear(self.dimV, self.dimV, bias = False)self.sqrtD = 1 / sqrt(dimQ)def forward(self, x):batch, n, dim = x.shapeassert(dim == self.dimQ)Q = self.linerQ(x)K = self.linerK(x)V = self.linerV(x)dist = torch.bmm(Q, K.transpose(1, 2)) * self.sqrtDW = torch.softmax(dist, dim = -1)Output = torch.bmm(W, V)return Outputif __name__ == "__main__":x = torch.tensor([[[1,2,3],[2,3,4],[3,4,5],[4,5,6]],[[1,2,3],[2,3,4],[3,4,5],[4,5,6]]], dtype = torch.float)print(x.shape)saModel = SA(3, 3, 3)Output = saModel(x)print(Output)

https://zhuanlan.zhihu.com/p/338817680
https://blog.csdn.net/weixin_44750512/article/details/124244915
https://blog.csdn.net/qq_40178291/article/details/100302375