以下是一个使用Transformer模型进行机器翻译的示例代码:
```python
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结束时,打印出平均损失。
请注意,上述代码仅为示例,需要根据实际情况进行修改和调整。