Transformer 注意力模块部分

1. Self_Attention

作用：

• 计算最基础的注意力分数和注意力输出。

• 输入是 Q、K、V，输出是加权后的结果 out 和注意力权重 attn。

它做了什么：

1. 用 Q @ K^T 计算相似度分数。

2. 用 sqrt(d_k) 做缩放，防止分数过大。

3. 如果传入 mask，就把不能看的位置设成负无穷。

4. 对分数做 softmax，得到注意力权重。

5. 用注意力权重对 V 做加权求和。

这个模块的核心意义：

• 决定“当前 token 应该重点看谁”。

#自注意力机制
class Self_Attention(nn.Module):
    def __init__(self,dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout) # 防止过拟合
        self.softmax = nn.Softmax(dim=-1) #转换为得分概率，归一化

    def forward(self,Q,K,V,mask=None):
        # x：batch，seq_len,d_model
        # batch:输入句子数,seq_len:一个句子的token数量，d_model;embedding向量的维度
        # Q。query向量 维度：batch，heads，seq_len_q,d_k
        # K。key向量 维度：batch，heads，seq_len_k,d_k
        # V。val向量 维度：batch，heads，seq_len_v,d_v
        # mask 掩码注意力机制需要，用来掩盖需要猜测的关系
        d_k = Q.size(-1)#Q,K的d_k相等
        scores = torch.matmul(Q,K.transpose(-2,-1)) / math.sqrt(d_k)

        if mask is not None:#如果mask不是None，为掩码注意力机制
            #掩码注意力机制，通过将mask==0（被屏蔽）的位置改为inf（负无穷），让需要猜测的位置进行屏蔽
            scores = scores.masked_fill(mask==0, float('-inf'))

        attn = self.softmax(scores)

        attn = self.dropout(attn)#再次防止过拟合

        out = torch.matmul(attn,V)

        return out,attn

2. MultiHeadAttention

作用：

• 把单头注意力扩展成多头注意力。

• 让模型能从多个角度同时学习 token 之间的关系。

它做了什么：

1. 用三个线性层把输入映射成 Q、K、V。

2. 把 d_model 拆成多个头，每个头单独做注意力。

3. 调用 Self_Attention 得到每个头的结果。

4. 把多个头拼接回来。

5. 再经过线性层、dropout、残差连接和 LayerNorm。

这个模块的核心意义：

• 不同头可以分别关注不同类型的信息，比如局部关系、长距离关系、语义对应关系。

#多头注意力机制
class MultiHeadAttention(nn.Module):
    def __init__(self,
                 d_model,
                 n_heads,
                 dropout=0.1):
        super().__init__()
        # d_model 需要能被n_head整除
        # d_model 为embedding的维度 512
        # n_head  为多头注意力机制的头数 64
        assert d_model % n_heads == 0

        self.d_k = d_model // n_heads

        self.n_heads = n_heads

        #输入映射到 QKV三个向量
        self.W_Q = nn.Linear(d_model,d_model)

        self.W_K = nn.Linear(d_model,d_model)

        self.W_V = nn.Linear(d_model,d_model)

        self.fc = nn.Linear(d_model,d_model)#拼接后在映射回原来的d_model

        self.attention = Self_Attention(dropout)

        self.dropout = nn.Dropout(dropout) #防止过拟合

        self.norm = nn.LayerNorm(d_model) #残差后的层归一化


    def forward(self,Q,K,V,mask=None):
        batch_size = Q.size(0)
        q = Q
        # q的维度 batch，seq_len,d_model -> batch,seq_len,n_heads,d_k ->batch,n_heads,seq_len,d_k
        #  为了让每个注意力头独立处理整个序列 ， 方便后续计算注意力权重
        Q = self.W_Q(Q).view(batch_size,-1,self.n_heads,self.d_k).transpose(1, 2)

        K = self.W_K(K).view(batch_size,-1,self.n_heads,self.d_k).transpose(1, 2)

        V = self.W_V(V).view(batch_size,-1,self.n_heads,self.d_k).transpose(1, 2)

        # 计算注意力
        out, attn = self.attention(Q,K,V,mask)

        out = out.transpose(1, 2).contiguous().view(batch_size,-1,self.d_k*self.n_heads)
        out = self.fc(out)
        out = self.dropout(out)
        return self.norm(out + q),attn