再写一个tansformer的完整代码

好的，这里是一个简单的 Transformer 的完整代码示例，用 PyTorch 实现。

首先，我们需要导入必要的库：

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

然后，我们定义一个 Transformer 模型类：

class Transformer(nn.Module):
    def __init__(self, input_dim, output_dim, d_model, n_heads, n_layers, dropout):
        super().__init__()
        self.input_embedding = nn.Linear(input_dim, d_model)
        self.output_embedding = nn.Linear(d_model, output_dim)
        self.positional_encoding = PositionalEncoding(d_model, dropout)
        self.encoder_layers = nn.ModuleList([
            EncoderLayer(d_model, n_heads, dropout) for _ in range(n_layers)
        ])
        self.decoder_layers = nn.ModuleList([
            DecoderLayer(d_model, n_heads, dropout) for _ in range(n_layers)
        ])
        self.output_projection = nn.Linear(d_model, output_dim)

    def forward(self, src, trg, src_mask=None, trg_mask=None):
        src = self.input_embedding(src)
        trg = self.input_embedding(trg)
        src = self.positional_encoding(src)
        trg = self.positional_encoding(trg)
        for layer in self.encoder_layers:
            src = layer(src, src_mask)
        for layer in self.decoder_layers:
            trg = layer(trg, src, trg_mask, src_mask)
        output = self.output_projection(trg)
        output = self.output_embedding(output)
        return output

这里的 `Transformer` 类包含了输入和输出的嵌入层，位置编码器，编码器和解码器层，以及最后的输出投影层。

接下来，我们需要定义位置编码器、编码器层和解码器层：

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, 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 EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, dropout):
        super().__init__()
        self.self_attention = MultiHeadAttention(d_model, n_heads)
        self.feedforward = FeedForward(d_model, dropout)

    def forward(self, x, mask):
        x = self.self_attention(x, x, x, mask)
        x = self.feedforward(x)
        return x


class DecoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, dropout):
        super().__init__()
        self.self_attention = MultiHeadAttention(d_model, n_heads)
        self.encoder_attention = MultiHeadAttention(d_model, n_heads)
        self.feedforward = FeedForward(d_model, dropout)

    def forward(self, x, encoder_output, trg_mask, src