核心原理：输入x(维度为Batch_size*seq_len *hidden_size)

一共有heads_num个头，每个头对都有对应的权重矩阵，如（维度为，也因此hidden_size必须可被head_num整除,否则会报错)

将权重矩阵拼接起来（维度为)，例如

这样

再按头拆分， 进行接下来的运算：也即

最后输出

import torch
import torch.nn as nn

#hidden_size:隐藏层维度
#head_num:注意力头数
#必须保证hidden_size可被head_num整除,否则会报错

class MHA(nn.Module):
    def __init__(self, hidden_size, heads_num):
        super().__init__() #python2环境则需改为super(MHA, self).__init__()
        self.hidden_size = hidden_size
        self.heads_num = heads_num
        self.head_size = hidden_size//heads_num

        # 线性变换矩阵
        # W_q(x)相当于x乘以[W_q_0,W_q_1,...,W_q_(heads_num-1)]矩阵,得到每个头的q值拼接后的结果
        # 同理W_k(x)和W_v(x)也相当于x乘以[W_k_0,W_k_1,...,W_k_(heads_num-1)]和[W_v_0,W_v_1,...,W_v_(heads_num-1)]矩阵,
        # 得到每个头的k值和v值拼接后的结果
        self.W_q = nn.Linear(hidden_size,hidden_size)
        self.W_k = nn.Linear(hidden_size,hidden_size)
        self.W_v = nn.Linear(hidden_size,hidden_size)

        self.W_o = nn.Linear(hidden_size,hidden_size)
    # x:输入序列，[batch_size, seq_len, hidden_size]
    def forward(self, x):
        q = self.W_q(x) # [batch_size, seq_len, hidden_size]
        k = self.W_k(x) # [batch_size, seq_len, hidden_size]
        v = self.W_v(x) # [batch_size, seq_len, hidden_size]

        # 将q、k、v分成多个头，每个头的维度为head_size
        q = q.view(q.size(0), q.size(1), self.heads_num, self.head_size) # [batch_size, heads_num, seq_len, head_size]
        k = k.view(k.size(0), k.size(1), self.heads_num, self.head_size) # [batch_size, heads_num, seq_len, head_size]
        v = v.view(v.size(0), v.size(1), self.heads_num, self.head_size) # [batch_size, heads_num, seq_len, head_size]

        # 翻转heads_num, seq_len维度，使得后续计算每个头更方便
        q = q.transpose(1,2) # [batch_size, heads_num, seq_len, head_size]
        k = k.transpose(1,2) # [batch_size, heads_num, seq_len, head_size]
        v = v.transpose(1,2) # [batch_size, heads_num, seq_len, head_size]

        # torch.matmul(q, k.transpose(-2,-1))相当于每个头的q值矩阵和k值矩阵的转置相乘得到每个头的注意力权重矩阵
        qk = torch.matmul(q, k.transpose(-2,-1))/torch.sqrt(torch.tensor(self.head_size, dtype = torch.float32))# [batch_size, heads_num, seq_len, seq_len]
        # 对每个头的注意力权重矩阵进行softmax归一化
        attn = torch.softmax(qk, dim=-1) # [batch_size, heads_num, seq_len, seq_len]

        # 将每个头的注意力权重矩阵和v值矩阵相乘得到每个头的输出矩阵
        out = torch.matmul(attn, v) # [batch_size, heads_num, seq_len, head_size]
        # 将每个头的输出矩阵拼接起来得到最终的输出矩阵
        # .contiguous() 时，PyTorch 会按照张量的逻辑形状，重新在内存中排列元素；返回一个新的、内存连续的张量
        # 如果省略 .contiguous()，直接执行 view 会抛出报错：
        # 核心原因：view 要求张量在内存中是连续的，否则无法按指定形状重塑。
        out = out.transpose(1,2).contiguous().view(out.size(0),out.size(2), self.hidden_size) # [batch_size, seq_len, hidden_size])
        # 最后通过线性变换矩阵W_o得到最终的输出
        out = self.W_o(out) # [batch_size, seq_len, hidden_size]

        return out







    

纯享版：

import torch
import torch.nn as nn

class MHA(nn.Module):
    def __init__(self, hidden_size,heads_num):
        super().__init__()
        self.hidden_size = hidden_size
        self.heads_num = heads_num
        self.head_size = hidden_size//heads_num

        self.W_q = nn.Linear(hidden_size,hidden_size)
        self.W_k = nn.Linear(hidden_size,hidden_size)
        self.W_v = nn.Linear(hidden_size,hidden_size)

        self.W_o = nn.Linear(hidden_size,hidden_size)

    def forward(self, x):
        q = self.W_q(x)
        k = self.W_k(x)
        v = self.W_v(x)

        q = q.view(q.size(0),q.size(1),self.heads_num,self.head_size).transpose(1,2)
        k = k.view(k.size(0),k.size(1),self.heads_num,self.head_size).transpose(1,2)
        v = v.view(v.size(0),v.size(1),self.heads_num,self.head_size).transpose(1,2)

        qk = torch.matmul(q,k.transpose(-2,-1))/torch.sqrt(torch.tensor(self.head_size,dtype = torch.float32))
        attn = torch.softmax(qk,dim=-1)
        out = torch.matmul(attn,v)

        out = out.transpose(1,2).contiguous().view(out.size(0),out.size(2),self.hidden_size)
        out = self.W_o(out)
        return out