importtorchimporttorch.nnasnnimportmathclassMultiheadAttention(nn.Module):# 注意: Module 首字母大写def__init__(self,d_model,n_head): 多头注意力机制 参数: d_model: 输入特征维度 n_head: 注意力头数量 super(MultiheadAttention,self).__init__()self.d_modeld_model self.n_headn_head self.n_dimd_model//n_head# 验证维度是否可整除assertd_model%n_head0,d_model must be divisible by n_head# 线性变换层self.q_weightnn.Linear(d_model,d_model)self.k_weightnn.Linear(d_model,d_model)self.v_weightnn.Linear(d_model,d_model)self.out_projnn.Linear(d_model,d_model)defforward(self,x): 前向传播 参数: x: 输入张量 [batch_size, seq_len, d_model] 返回: 输出张量 [batch_size, seq_len, d_model] bs,seq_len,dimx.size()# 线性变换qself.q_weight(x)# [bs, seq_len, d_model]kself.k_weight(x)# [bs, seq_len, d_model]vself.v_weight(x)# [bs, seq_len, d_model]# 分割多头 - 更稳健的实现# 首先view到 [bs, seq_len, n_head, n_dim]qq.view(bs,seq_len,self.n_head,self.n_dim)kk.view(bs,seq_len,self.n_head,self.n_dim)vv.view(bs,seq_len,self.n_head,self.n_dim)# 转置为 [bs, n_head, seq_len, n_dim]qq.transpose(1,2)kk.transpose(1,2)vv.transpose(1,2)# 计算注意力分数# [bs, n_head, seq_len, n_dim] [bs, n_head, n_dim, seq_len]# - [bs, n_head, seq_len, seq_len]attn_scorestorch.matmul(q,k.transpose(-2,-1))/math.sqrt(self.n_dim)# 计算注意力权重attn_weightstorch.softmax(attn_scores,dim-1)# 应用注意力权重到值向量# [bs, n_head, seq_len, seq_len] [bs, n_head, seq_len, n_dim]# - [bs, n_head, seq_len, n_dim]attn_outputtorch.matmul(attn_weights,v)# 合并多头# 先转置为 [bs, seq_len, n_head, n_dim]attn_outputattn_output.transpose(1,2)# 合并最后两个维度 [bs, seq_len, n_head * n_dim] [bs, seq_len, d_model]attn_outputattn_output.contiguous().view(bs,seq_len,self.d_model)# 输出投影outputself.out_proj(attn_output)returnoutput# 测试代码if__name____main__:# 参数设置d_model64n_head8batch_size2seq_len10# 创建多头注意力层mhaMultiheadAttention(d_model,n_head)# 创建随机输入xtorch.randn(batch_size,seq_len,d_model)# 测试前向传播outputmha(x)# 验证输出形状expected_shape(batch_size,seq_len,d_model)actual_shapeoutput.shapeprint(f输入形状:{x.shape})print(f输出形状:{actual_shape})# 验证是否保留原始维度assertactual_shapeexpected_shape,\f形状错误! 期望{expected_shape}, 实际{actual_shape}print(测试通过输出与输入形状一致保留原始维度。)