bert后接attention代码

下面是BERT模型中的Attention代码示例：

class MultiHeadAttention(nn.Module):
    def __init__(self, hidden_size, num_heads, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads
        self.dropout = nn.Dropout(dropout)
        
        # 定义线性变换矩阵
        self.q_linear = nn.Linear(hidden_size, hidden_size)
        self.k_linear = nn.Linear(hidden_size, hidden_size)
        self.v_linear = nn.Linear(hidden_size, hidden_size)
        self.out_linear = nn.Linear(hidden_size, hidden_size)
        
    def split_heads(self, x, batch_size):
        """
        将输入的x按照头数进行划分，返回划分后的张量
        """
        x = x.view(batch_size, -1, self.num_heads, self.head_size)
        return x.transpose(1, 2)
        
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # 线性变换
        q = self.q_linear(query)
        k = self.k_linear(key)
        v = self.v_linear(value)
        
        # 按照头数进行划分
        q = self.split_heads(q, batch_size)
        k = self.split_heads(k, batch_size)
        v = self.split_heads(v, batch_size)
        
        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_size)
        
        # mask
        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # softmax
        attention = nn.Softmax(dim=-1)(scores)
        
        # dropout
        attention = self.dropout(attention)
        
        # 加权求和
        context = torch.matmul(attention, v)
        
        # 合并多头
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.hidden_size)
        
        # 线性变换
        output = self.out_linear(context)
        
        return output

该代码实现了BERT中的多头自注意力机制。其中，输入的query、key和value都是形状为[batch_size, seq_len, hidden_size]的三维张量。在MultiHeadAttention类的forward方法中，首先对query、key和value进行线性变换，然后按照头数进行划分。接着，计算注意力分数并进行softmax操作，然后进行dropout和加权求和。最后，将多头合并，再进行一次线性变换，输出形状为[batch_size, seq_len, hidden_size]的张量。