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

采用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_dimconfig.hidden_size self.num_attention_headsconfig.num_attention_heads self.head_dimself.emb_dim//self.num_attention_heads self.q_projnn.Linear(self.emb_dim,self.emb_dim) self.k_projnn.Linear(self.emb_dim,self.emb_dim) self.v_projnn.Linear(self.emb_dim,self.emb_dim) self.o_projnn.Linear(self.emb_dim,self.emb_dim) def forward(self,x,attn_maskNone): qself.q_proj(x) kself.k_proj(x) vself.v_proj(x) batch_size,seq_len,_q.size() qq.view(batch_size,seq_len,self.num_attention_heads,self.head_dim).transpose(1,2) kk.view(batch_size,seq_len,self.num_attention_heads,self.head_dim).transpose(1,2) vv.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_scoretorch.matmul(q,k.transpose(-1,-2)) # batch_size,self.num_attention_heads,seq_len,seq_len dkk.size(-1) attn_scoreattn_score/math.sqrt(dk) if attn_mask is not None: attn_scoreattn_score.masked_fill(attn_mask0,-float(inf)) weightsF.softmax(attn_score,dim-1) outputtorch.matmul(weights,v) # batch_size,self.num_attention_heads,seq_len,self.head_dim outputoutput.transpose(1,2).contiguous().view(batch_size,seq_len,self.emb_dim) outputself.o_proj(output) return output注意点contiguous()某些操作如transpose,permute,view,reshape等会改变张量的视图view而不复制数据导致逻辑顺序与内存物理顺序不一致即变成non-contiguous。而有些 PyTorch 操作如view()要求输入必须是 contiguous 的否则会报错。