大模型相关代码1 -- 注意力机制【手撕基础模型】

• 看见网上有要手撕注意力的面经，自己开始写，以为很简单，实际上自己菜的要死，遂写本博客。

公式

$$Attention(Q,K,V)=softmax(\frac{Q{K}^t}{\sqrt{d}})V$$

代码

def myattention(q, k, v, d_k, mask=None, dropout=None):
	# 首先计算注意力得分
    score = torch.matmul(q,k.transpose(-2,-1)) / math.sqrt(d_k)
    # 处理mask信息，如果被mask就变成一个很小的数
    if mask is not None:
        mask = mask.unsqueeze(1)
        score = score.masked_fill(mask == 0,-1e9)
    # 对注意力得分进行softmax
    score = F.softmax(score, dim=-1) # 沿着最后一个维度计算 其他维度不变
    # 处理dropout
    if dropout is not None:
        score = dropout(score)
    # 点乘value矩阵
    ans = torch.matmul(score,v)
    return ans