好的，下面是一个基于PyTorch框架实现的Transformer模型的实战案例代码，以机器翻译为例：
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
# 定义超参数
MAX_LENGTH = 10 # 最大句子长度
BATCH_SIZE = 64 # 批量大小
EMBEDDING_DIM = 256 # 词向量维度
HIDDEN_DIM = 512 # 隐藏层维度
N_LAYERS = 6 # Transformer层数
N_HEADS = 8 # 注意力头数
DROPOUT = 0.1 # Dropout概率
LEARNING_RATE = 0.0005 # 学习率
N_EPOCHS = 10 # 训练轮数
# 定义数据集
class TranslationDataset(Dataset):
def __init__(self, path):
self.pairs = []
with open(path, 'r', encoding='utf-8') as f:
for line in f:
src, tgt = line.strip().split('\t')
self.pairs.append((src, tgt))
def __len__(self):
return len(self.pairs)
def __getitem__(self, index):
src, tgt = self.pairs[index]
src_tokens = [int(x) for x in src.split()]
tgt_tokens = [int(x) for x in tgt.split()]
src_len = len(src_tokens)
tgt_len = len(tgt_tokens)
src_padding = [0] * (MAX_LENGTH - src_len)
tgt_padding = [0] * (MAX_LENGTH - tgt_len)
src_tokens += src_padding
tgt_tokens += tgt_padding
return (torch.LongTensor(src_tokens), torch.LongTensor(tgt_tokens))
# 定义Transformer模型
class Transformer(nn.Module):
def __init__(self, input_dim, output_dim, embedding_dim, hidden_dim, n_layers, n_heads, dropout):
super().__init__()
self.input_embedding = nn.Embedding(input_dim, embedding_dim)
self.output_embedding = nn.Embedding(output_dim, embedding_dim)
self.encoder = Encoder(embedding_dim, hidden_dim, n_layers, n_heads, dropout)
self.decoder = Decoder(embedding_dim, hidden_dim, n_layers, n_heads, dropout)
self.output_projection = nn.Linear(hidden_dim, output_dim)
def forward(self, src, tgt):
src_embedded = self.input_embedding(src)
tgt_embedded = self.output_embedding(tgt)
encoder_output, encoder_attention = self.encoder(src_embedded)
decoder_output, decoder_attention = self.decoder(tgt_embedded, encoder_output)
output = self.output_projection(decoder_output)
return output, encoder_attention, decoder_attention
class Encoder(nn.Module):
def __init__(self, embedding_dim, hidden_dim, n_layers, n_heads, dropout):
super().__init__()
self.layers = nn.ModuleList([EncoderLayer(embedding_dim, hidden_dim, n_heads, dropout) for _ in range(n_layers)])
self.dropout = nn.Dropout(dropout)
def forward(self, x):
attention_weights = []
for layer in self.layers:
x, attention = layer(x)
attention_weights.append(attention)
return x, torch.stack(attention_weights)
class EncoderLayer(nn.Module):
def __init__(self, embedding_dim, hidden_dim, n_heads, dropout):
super().__init__()
self.self_attention = MultiHeadAttention(embedding_dim, n_heads)
self.layer_norm1 = nn.LayerNorm(embedding_dim)
self.positionwise_feedforward = PositionwiseFeedforward(embedding_dim, hidden_dim, dropout)
self.layer_norm2 = nn.LayerNorm(embedding_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
residual = x
x, attention = self.self_attention(x, x, x)
x = self.layer_norm1(residual + self.dropout(x))
residual = x
x = self.positionwise_feedforward(x)
x = self.layer_norm2(residual + self.dropout(x))
return x, attention
class Decoder(nn.Module):
def __init__(self, embedding_dim, hidden_dim, n_layers, n_heads, dropout):
super().__init__()
self.layers = nn.ModuleList([DecoderLayer(embedding_dim, hidden_dim, n_heads, dropout) for _ in range(n_layers)])
self.dropout = nn.Dropout(dropout)
def forward(self, x, encoder_output):
attention_weights = []
for layer in self.layers:
x, attention = layer(x, encoder_output)
attention_weights.append(attention)
return x, torch.stack(attention_weights)
class DecoderLayer(nn.Module):
def __init__(self, embedding_dim, hidden_dim, n_heads, dropout):
super().__init__()
self.self_attention = MultiHeadAttention(embedding_dim, n_heads)
self.layer_norm1 = nn.LayerNorm(embedding_dim)
self.encoder_attention = MultiHeadAttention(embedding_dim, n_heads)
self.layer_norm2 = nn.LayerNorm(embedding_dim)
self.positionwise_feedforward = PositionwiseFeedforward(embedding_dim, hidden_dim, dropout)
self.layer_norm3 = nn.LayerNorm(embedding_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x, encoder_output):
residual = x
x, self_attention = self.self_attention(x, x, x)
x = self.layer_norm1(residual + self.dropout(x))
residual = x
x, encoder_attention = self.encoder_attention(x, encoder_output, encoder_output)
x = self.layer_norm2(residual + self.dropout(x))
residual = x
x = self.positionwise_feedforward(x)
x = self.layer_norm3(residual + self.dropout(x))
return x, encoder_attention
class MultiHeadAttention(nn.Module):
def __init__(self, embedding_dim, n_heads):
super().__init__()
self.embedding_dim = embedding_dim
self.n_heads = n_heads
self.head_dim = embedding_dim // n_heads
self.q_linear = nn.Linear(embedding_dim, embedding_dim)
self.k_linear = nn.Linear(embedding_dim, embedding_dim)
self.v_linear = nn.Linear(embedding_dim, embedding_dim)
self.out_linear = nn.Linear(embedding_dim, embedding_dim)
def forward(self, query, key, value):
batch_size = query.size(0)
Q = self.q_linear(query).view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)
K = self.k_linear(key).view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)
V = self.v_linear(value).view(batch_size, -1, self.n_heads, self.head_dim).transpose(1, 2)
attention_weights = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
attention_weights = F.softmax(attention_weights, dim=-1)
output = torch.matmul(self.dropout(attention_weights), V)
output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.embedding_dim)
output = self.out_linear(output)
return output, attention_weights
class PositionwiseFeedforward(nn.Module):
def __init__(self, embedding_dim, hidden_dim, dropout):
super().__init__()
self.fc1 = nn.Linear(embedding_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, embedding_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = self.fc2(x)
return x
# 定义训练函数
def train(model, iterator, optimizer, criterion):
model.train()
epoch_loss = 0
for src, tgt in iterator:
optimizer.zero_grad()
output, _, _ = model(src, tgt[:, :-1])
output_dim = output.shape[-1]
output = output.contiguous().view(-1, output_dim)
tgt = tgt[:, 1:].contiguous().view(-1)
loss = criterion(output, tgt)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
return epoch_loss / len(iterator)
# 定义评估函数
def evaluate(model, iterator, criterion):
model.eval()
epoch_loss = 0
with torch.no_grad():
for src, tgt in iterator:
output, _, _ = model(src, tgt[:, :-1])
output_dim = output.shape[-1]
output = output.contiguous().view(-1, output_dim)
tgt = tgt[:, 1:].contiguous().view(-1)
loss = criterion(output, tgt)
epoch_loss += loss.item()
return epoch_loss / len(iterator)
# 加载数据集
train_dataset = TranslationDataset('train.txt')
valid_dataset = TranslationDataset('valid.txt')
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE)
# 初始化模型和优化器
model = Transformer(input_dim=10000, output_dim=10000, embedding_dim=EMBEDDING_DIM, hidden_dim=HIDDEN_DIM, n_layers=N_LAYERS, n_heads=N_HEADS, dropout=DROPOUT)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = nn.CrossEntropyLoss(ignore_index=0)
# 训练模型
for epoch in range(N_EPOCHS):
train_loss = train(model, train_loader, optimizer, criterion)
valid_loss = evaluate(model, valid_loader, criterion)
print(f'Epoch: {epoch+1:02} | Train Loss: {train_loss:.3f} | Valid Loss: {valid_loss:.3f}')
```
上述代码实现了一个简单的机器翻译模型，使用了PyTorch框架和Transformer架构。具体而言，它使用了一个包含6个Encoder层和6个Decoder层的Transformer模型，其中每个层都包含了自注意力和多头注意力机制。该模型使用了Adam优化器和交叉熵损失函数进行训练。
955
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
