大模型手撕代码（一）多头注意力机制

采用bert的config用于后续写代码时使用

{ "architectures": [ "BertForMaskedLM" ], "attention_probs_dropout_prob": 0.1, "directionality": "bidi", "hidden_act": "gelu", "hidden_dropout_prob": 0.1, "hidden_size": 768, "initializer_range": 0.02, "intermediate_size": 3072, "layer_norm_eps": 1e-12, "max_position_embeddings": 512, "model_type": "bert", "num_attention_heads": 12, "num_hidden_layers": 12, "pad_token_id": 0, "pooler_fc_size": 768, "pooler_num_attention_heads": 12, "pooler_num_fc_layers": 3, "pooler_size_per_head": 128, "pooler_type": "first_token_transform", "type_vocab_size": 2, "vocab_size": 21128, "return_dict": false, "num_labels":18 }

根据config中的参数写多头注意力机制的代码

import torch import torch.nn as nn import torch.nn.functional as F import math class mha(nn.Module): def __init__(self,config): super().__init__() self.emb_dim=config.hidden_size self.num_attention_heads=config.num_attention_heads self.head_dim=self.emb_dim//self.num_attention_heads self.q_proj=nn.Linear(self.emb_dim,self.emb_dim) self.k_proj=nn.Linear(self.emb_dim,self.emb_dim) self.v_proj=nn.Linear(self.emb_dim,self.emb_dim) self.o_proj=nn.Linear(self.emb_dim,self.emb_dim) def forward(self,x,attn_mask=None): q=self.q_proj(x) k=self.k_proj(x) v=self.v_proj(x) batch_size,seq_len,_=q.size() q=q.view(batch_size,seq_len,self.num_attention_heads,self.head_dim).transpose(1,2) k=k.view(batch_size,seq_len,self.num_attention_heads,self.head_dim).transpose(1,2) v=v.view(batch_size,seq_len,self.num_attention_heads,self.head_dim).transpose(1,2) # batch_size,self.num_attention_heads,seq_len,self.head_dim attn_score=torch.matmul(q,k.transpose(-1,-2)) # batch_size,self.num_attention_heads,seq_len,seq_len dk=k.size(-1) attn_score=attn_score/math.sqrt(dk) if attn_mask is not None: attn_score=attn_score.masked_fill(attn_mask==0,-float('inf')) weights=F.softmax(attn_score,dim=-1) output=torch.matmul(weights,v) # batch_size,self.num_attention_heads,seq_len,self.head_dim output=output.transpose(1,2).contiguous().view(batch_size,seq_len,self.emb_dim) output=self.o_proj(output) return output

注意点：contiguous()

某些操作（如transpose,permute,view,reshape等）会改变张量的视图（view），而不复制数据，导致逻辑顺序与内存物理顺序不一致，即变成non-contiguous。而有些 PyTorch 操作（如view()）要求输入必须是 contiguous 的，否则会报错。