企业官网建设流程全解析

importtorchimporttorch.nnasnnimportmathclassMultiheadAttention(nn.Module):# 注意: Module 首字母大写def__init__(self,d_model,n_head):""" 多头注意力机制 参数: d_model: 输入特征维度 n_head: 注意力头数量 """super(MultiheadAttention,self).__init__()self.d_model=d_model self.n_head=n_head self.n_dim=d_model//n_head# 验证维度是否可整除assertd_model%n_head==0,"d_model must be divisible by n_head"# 线性变换层self.q_weight=nn.Linear(d_model,d_model)self.k_weight=nn.Linear(d_model,d_model)self.v_weight=nn.Linear(d_model,d_model)self.out_proj=nn.Linear(d_model,d_model)defforward(self,x):""" 前向传播 参数: x: 输入张量 [batch_size, seq_len, d_model] 返回: 输出张量 [batch_size, seq_len, d_model] """bs,seq_len,dim=x.size()# 线性变换q=self.q_weight(x)# [bs, seq_len, d_model]k=self.k_weight(x)# [bs, seq_len, d_model]v=self.v_weight(x)# [bs, seq_len, d_model]# 分割多头 - 更稳健的实现# 首先view到 [bs, seq_len, n_head, n_dim]q=q.view(bs,seq_len,self.n_head,self.n_dim)k=k.view(bs,seq_len,self.n_head,self.n_dim)v=v.view(bs,seq_len,self.n_head,self.n_dim)# 转置为 [bs, n_head, seq_len, n_dim]q=q.transpose(1,2)k=k.transpose(1,2)v=v.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_scores=torch.matmul(q,k.transpose(-2,-1))/math.sqrt(self.n_dim)# 计算注意力权重attn_weights=torch.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_output=torch.matmul(attn_weights,v)# 合并多头# 先转置为 [bs, seq_len, n_head, n_dim]attn_output=attn_output.transpose(1,2)# 合并最后两个维度 [bs, seq_len, n_head * n_dim] = [bs, seq_len, d_model]attn_output=attn_output.contiguous().view(bs,seq_len,self.d_model)# 输出投影output=self.out_proj(attn_output)returnoutput# 测试代码if__name__=="__main__":# 参数设置d_model=64n_head=8batch_size=2seq_len=10# 创建多头注意力层mha=MultiheadAttention(d_model,n_head)# 创建随机输入x=torch.randn(batch_size,seq_len,d_model)# 测试前向传播output=mha(x)# 验证输出形状expected_shape=(batch_size,seq_len,d_model)actual_shape=output.shapeprint(f"输入形状:{x.shape}")print(f"输出形状:{actual_shape}")# 验证是否保留原始维度assertactual_shape==expected_shape,\f"形状错误! 期望{expected_shape}, 实际{actual_shape}"print("测试通过！输出与输入形状一致，保留原始维度。")