用pytorch的MultiheadAttention

MultiheadAttention是PyTorch提供的一个模块，可以用来实现多头注意力机制。它可以将一个输入序列分成多个头，每个头进行注意力计算，并将多个头的输出进行拼接和线性变换，得到最终的输出。使用MultiheadAttention需要以下步骤：

1.定义MultiheadAttention模块：

import torch.nn as nn

class MultiheadAttention(nn.Module):
    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        super(MultiheadAttention, self).__init__()

        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, q, k, v, mask=None):
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, _ = q.size()
        sz_b, len_k, _ = k.size()
        sz_b, len_v, _ = v.size()

        residual = q

        # project inputs to multi-heads
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        # transpose to prepare for softmax
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        # calculate attention scores
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # apply softmax to get attention probabilities
        attn = nn.Softmax(dim=-1)(scores)
        attn = self.dropout(attn)

        # apply attention probabilities to values
        output = torch.matmul(attn, v)

        # concatenate multi-heads and apply output projection
        output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        output = self.fc(output)

        # add residual and apply layer normalization
        output = self.layer_norm(output + residual)

        return output, attn

2.使用MultiheadAttention模块：

import torch

# define inputs
q = torch.randn(2, 4, 128)  # query
k = torch.randn(2, 6, 128)  # key
v = torch.randn(2, 6, 128)  # value

# define MultiheadAttention module
attn = MultiheadAttention(n_head=8, d_model=128, d_k=16, d_v=16)

# apply MultiheadAttention
output, attn_weights = attn(q, k, v)

# output shape: [2, 4, 128]
# attn_weights shape: [2, 8, 4, 6]