使用PyTorch或TensorFlow等深度学习框架实现Transformer的代码

以下是使用PyTorch实现Transformer的代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, heads, d_model):
        super().__init__()
        self.heads = heads
        self.d_model = d_model
        self.d_head = d_model // heads

        self.q_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)

        self.out = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):
        bs = q.size(0)

        # Linear projections
        k = self.k_linear(k).view(bs, -1, self.heads, self.d_head)
        q = self.q_linear(q).view(bs, -1, self.heads, self.d_head)
        v = self.v_linear(v).view(bs, -1, self.heads, self.d_head)

        # Transpose and dot product attention
        k = k.transpose(1,2)
        q = q.transpose(1,2)
        v = v.transpose(1,2)

        scores = torch.matmul(q, k.transpose(-2, -1)) /  math.sqrt(self.d_head)

        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)

        scores = F.softmax(scores, dim=-1)

        # Output attention
        output = torch.matmul(scores, v)

        # Concatenate and linear projection
        output = output.transpose(1,2).contiguous().view(bs, -1, self.d_model)
        return self.out(output)

class PositionwiseFeedforward(nn.Module):
    def __init__(self, d_model, d_ff=2048):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        x = self.linear1(x)
        x = F.relu(x)
        x = self.linear2(x)
        return x

class EncoderLayer(nn.Module):
    def __init__(self, d_model, heads, dropout=0.1):
        super().__init__()

        self.norm_1 = nn.LayerNorm(d_model)
        self.norm_2 = nn.LayerNorm(d_model)

        self.attn = MultiHeadAttention(heads, d_model)
        self.ff = PositionwiseFeedforward(d_model)

        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)

    def forward(self, x, mask):
        x2 = self.norm_1(x)
        x = x + self.dropout_1(self.attn(x2, x2, x2, mask))
        x2 = self.norm_2(x)
        x = x + self.dropout_2(self.ff(x2))
        return x

class TransformerEncoder(nn.Module):
    def __init__(self, input_dim, d_model, heads, num_layers):
        super().__init__()
        self.input_dim = input_dim
        self.d_model = d_model
        self.heads = heads
        self.num_layers = num_layers

        self.embedding = nn.Embedding(input_dim, d_model)
        self.pe = PositionalEncoder(d_model)
        self.layers = nn.ModuleList([EncoderLayer(d_model, heads) for _ in range(num_layers)])

    def forward(self, src_seq, src_mask):
        x = self.embedding(src_seq)
        x = self.pe(x)
        for i in range(self.num_layers):
            x = self.layers[i](x, src_mask)
        return x

class PositionalEncoder(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

class Transformer(nn.Module):
    def __init__(self, input_dim, output_dim, d_model, heads, num_layers, dropout=0.1):
        super().__init__()

        self.encoder = TransformerEncoder(input_dim, d_model, heads, num_layers)
        self.fc = nn.Linear(d_model, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src_seq, src_mask):
        x = self.encoder(src_seq, src_mask)
        x = x.mean(dim=1)
        x = self.fc(x)
        return x

以上是使用PyTorch实现Transformer的代码，其中包括了Transformer的编码器、多头自注意力机制、位置编码等模块。需要注意的是，该代码中使用了Layer Normalization进行层归一化。