transformer实例实现代码

以下是一个使用Transformer模型进行机器翻译的示例代码：

import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import Adam
from torch.utils.data import DataLoader, Dataset

class TranslationDataset(Dataset):

    def __init__(self, src_sentences, tgt_sentences):
        self.src_sentences = src_sentences
        self.tgt_sentences = tgt_sentences

    def __len__(self):
        return len(self.src_sentences)

    def __getitem__(self, idx):
        src_sentence = self.src_sentences[idx]
        tgt_sentence = self.tgt_sentences[idx]
        return src_sentence, tgt_sentence

class TransformerModel(nn.Module):

    def __init__(self, vocab_size, embedding_size, hidden_size, num_layers, num_heads, dropout_prob):
        super(TransformerModel, self).__init__()

        self.embedding_size = embedding_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.num_heads = num_heads

        self.embedding = nn.Embedding(vocab_size, embedding_size)
        self.positional_encoding = PositionalEncoding(embedding_size, dropout_prob)

        encoder_layer = nn.TransformerEncoderLayer(embedding_size, num_heads, hidden_size, dropout_prob)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)

        decoder_layer = nn.TransformerDecoderLayer(embedding_size, num_heads, hidden_size, dropout_prob)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers)

        self.output_layer = nn.Linear(embedding_size, vocab_size)

    def forward(self, src, tgt):
        src_embeddings = self.embedding(src)
        tgt_embeddings = self.embedding(tgt)

        src_embeddings = src_embeddings * (self.embedding_size ** 0.5)
        src_embeddings = self.positional_encoding(src_embeddings)

        tgt_embeddings = tgt_embeddings * (self.embedding_size ** 0.5)
        tgt_embeddings = self.positional_encoding(tgt_embeddings)

        memory = self.encoder(src_embeddings)
        output = self.decoder(tgt_embeddings, memory)

        output = self.output_layer(output)

        return F.log_softmax(output, dim=-1)

class PositionalEncoding(nn.Module):

    def __init__(self, embedding_size, dropout_prob, max_length=5000):
        super(PositionalEncoding, self).__init__()

        pe = torch.zeros(max_length, embedding_size)
        position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embedding_size, 2).float() * (-math.log(10000.0) / embedding_size))
        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)
        self.dropout = nn.Dropout(dropout_prob)

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

# 准备训练数据和目标数据
src_sentences = ["I am happy", "She is sad", "They are dancing"]
tgt_sentences = ["Je suis heureux", "Elle est triste", "Ils dansent"]

# 创建数据集和数据加载器
dataset = TranslationDataset(src_sentences, tgt_sentences)
dataloader = DataLoader(dataset, batch_size=3, shuffle=True)

# 定义模型和优化器
model = TransformerModel(vocab_size=10, embedding_size=32, hidden_size=64, num_layers=2, num_heads=4, dropout_prob=0.1)
optimizer = Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    total_loss = 0
    for src, tgt in dataloader:
        optimizer.zero_grad()
        src = torch.tensor(src)
        tgt = torch.tensor(tgt)
        output = model(src, tgt[:, :-1])
        loss = F.nll_loss(output.view(-1, vocab_size), tgt[:, 1:].contiguous().view(-1))
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(dataloader):.4f}")

上述代码中，我们创建了一个名为`TranslationDataset`的自定义数据集，用于存储源句子和目标句子。然后，我们使用`DataLoader`将数据加载到模型中进行训练。我们定义了一个名为`TransformerModel`的模型，并使用Adam优化器进行参数优化。在训练过程中，我们使用负对数似然损失函数(`nll_loss`)计算损失，并使用反向传播和梯度下降来更新模型的参数。每个epoch结束时，打印出平均损失。

请注意，上述代码仅为示例，需要根据实际情况进行修改和调整。