PyTorch 构建 Transformer 模型
Transformer 是现代机器学习中最强大的模型之一。
Transformer 模型是一种基于自注意力机制（Self-Attention） 的深度学习架构，它彻底改变了自然语言处理（NLP）领域，并成为现代深度学习模型（如 BERT、GPT 等）的基础。
Transformer 是现代 NLP 领域的核心架构，凭借其强大的长距离依赖建模能力和高效的并行计算优势，在语言翻译和文本摘要等任务中超越了传统的 长短期记忆 (LSTM) 网络。
如果你还不了解 Transformer，可以参考：Transformer 模型介绍。
使用 PyTorch 构建 Transformer 模型
构建 Transformer 模型的步骤如下：
1、导入必要的库和模块
导入 PyTorch 核心库、神经网络模块、优化器模块、数据处理工具，以及数学和对象复制模块，为定义模型架构、管理数据和训练过程提供支持。
import torch import torch.nn as nn import torch.optim as optim import torch.utils.data as data import math import copy
说明：
torch：PyTorch 的核心库，用于张量操作和自动求导。torch.nn：PyTorch 的神经网络模块，包含各种层和损失函数。torch.optim：优化算法模块，如 Adam、SGD 等。math：数学函数库，用于计算平方根等。copy：用于深度复制对象。
定义基本构建块：多头注意力、位置前馈网络、位置编码
多头注意力通过多个"注意力头"计算序列中每对位置之间的关系，能够捕捉输入序列的不同特征和模式。
MultiHeadAttention 类封装了 Transformer 模型中常用的多头注意力机制，负责将输入拆分成多个注意力头，对每个注意力头施加注意力，然后将结果组合起来，这样模型就可以在不同尺度上捕捉输入数据中的各种关系，提高模型的表达能力。
实例
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
self.d_model = d_model # 模型的维度
self.num_heads = num_heads # 注意力头的数量
self.d_k = d_model // num_heads # 每个头的维度
# 定义线性变换层
self.W_q = nn.Linear(d_model, d_model) # 查询变换
self.W_k = nn.Linear(d_model, d_model) # 键变换
self.W_v = nn.Linear(d_model, d_model) # 值变换
self.W_o = nn.Linear(d_model, d_model) # 输出变换
def scaled_dot_product_attention(self, Q, K, V, mask=None):
# 计算注意力分数
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
# 如果提供了掩码，则应用掩码
if mask is not None:
attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
# 使用 softmax 计算注意力权重
attn_probs = torch.softmax(attn_scores, dim=-1)
# 使用注意力权重对值进行加权求和
output = torch.matmul(attn_probs, V)
return output
def split_heads(self, x):
# 将输入张量分割成多个头
batch_size, seq_length, d_model = x.size()
return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)
def combine_heads(self, x):
# 将多个头合并回原始形状
batch_size, _, seq_length, d_k = x.size()
return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)
def forward(self, Q, K, V, mask=None):
# 对查询、键和值进行线性变换并分割成多个头
Q = self.split_heads(self.W_q(Q))
K = self.split_heads(self.W_k(K))
V = self.split_heads(self.W_v(V))
# 计算缩放点积注意力
attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
# 合并多头并应用输出变换
output = self.W_o(self.combine_heads(attn_output))
return output
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
self.d_model = d_model # 模型的维度
self.num_heads = num_heads # 注意力头的数量
self.d_k = d_model // num_heads # 每个头的维度
# 定义线性变换层
self.W_q = nn.Linear(d_model, d_model) # 查询变换
self.W_k = nn.Linear(d_model, d_model) # 键变换
self.W_v = nn.Linear(d_model, d_model) # 值变换
self.W_o = nn.Linear(d_model, d_model) # 输出变换
def scaled_dot_product_attention(self, Q, K, V, mask=None):
# 计算注意力分数
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
# 如果提供了掩码，则应用掩码
if mask is not None:
attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
# 使用 softmax 计算注意力权重
attn_probs = torch.softmax(attn_scores, dim=-1)
# 使用注意力权重对值进行加权求和
output = torch.matmul(attn_probs, V)
return output
def split_heads(self, x):
# 将输入张量分割成多个头
batch_size, seq_length, d_model = x.size()
return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)
def combine_heads(self, x):
# 将多个头合并回原始形状
batch_size, _, seq_length, d_k = x.size()
return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)
def forward(self, Q, K, V, mask=None):
# 对查询、键和值进行线性变换并分割成多个头
Q = self.split_heads(self.W_q(Q))
K = self.split_heads(self.W_k(K))
V = self.split_heads(self.W_v(V))
# 计算缩放点积注意力
attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
# 合并多头并应用输出变换
output = self.W_o(self.combine_heads(attn_output))
return output
说明：
多头注意力机制：将输入分割成多个头，每个头独立计算注意力，最后将结果合并。
缩放点积注意力：计算查询和键的点积，缩放后使用 softmax 计算注意力权重，最后对值进行加权求和。
掩码：用于屏蔽无效位置（如填充部分）。
位置前馈网络（Position-wise Feed-Forward Network）
实例
class PositionWiseFeedForward(nn.Module):
def __init__(self, d_model, d_ff):
super(PositionWiseFeedForward, self).__init__()
self.fc1 = nn.Linear(d_model, d_ff) # 第一层全连接
self.fc2 = nn.Linear(d_ff, d_model) # 第二层全连接
self.relu = nn.ReLU() # 激活函数
def forward(self, x):
# 前馈网络的计算
return self.fc2(self.relu(self.fc1(x)))
def __init__(self, d_model, d_ff):
super(PositionWiseFeedForward, self).__init__()
self.fc1 = nn.Linear(d_model, d_ff) # 第一层全连接
self.fc2 = nn.Linear(d_ff, d_model) # 第二层全连接
self.relu = nn.ReLU() # 激活函数
def forward(self, x):
# 前馈网络的计算
return self.fc2(self.relu(self.fc1(x)))
前馈网络：由两个全连接层和一个 ReLU 激活函数组成，用于进一步处理注意力机制的输出。
位置编码
位置编码用于注入输入序列中每个 token 的位置信息。
使用不同频率的正弦和余弦函数来生成位置编码。
实例
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_seq_length):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_seq_length, d_model) # 初始化位置编码矩阵
position = torch.arange(0, max_seq_length, 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) # 奇数位置使用余弦函数
self.register_buffer('pe', pe.unsqueeze(0)) # 注册为缓冲区
def forward(self, x):
# 将位置编码添加到输入中
return x + self.pe[:, :x.size(1)]
def __init__(self, d_model, max_seq_length):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_seq_length, d_model) # 初始化位置编码矩阵
position = torch.arange(0, max_seq_length, 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) # 奇数位置使用余弦函数
self.register_buffer('pe', pe.unsqueeze(0)) # 注册为缓冲区
def forward(self, x):
# 将位置编码添加到输入中
return x + self.pe[:, :x.size(1)]
构建编码器块（Encoder Layer）
编码器层：包含一个自注意力机制和一个前馈网络，每个子层后接残差连接和层归一化。
实例
class EncoderLayer(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads) # 自注意力机制
self.feed_forward = PositionWiseFeedForward(d_model, d_ff) # 前馈网络
self.norm1 = nn.LayerNorm(d_model) # 层归一化
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout) # Dropout
def forward(self, x, mask):
# 自注意力机制
attn_output = self.self_attn(x, x, x, mask)
x = self.norm1(x + self.dropout(attn_output)) # 残差连接和层归一化
# 前馈网络
ff_output = self.feed_forward(x)
x = self.norm2(x + self.dropout(ff_output)) # 残差连接和层归一化
return x
def __init__(self, d_model, num_heads, d_ff, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads) # 自注意力机制
self.feed_forward = PositionWiseFeedForward(d_model, d_ff) # 前馈网络
self.norm1 = nn.LayerNorm(d_model) # 层归一化
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout) # Dropout
def forward(self, x, mask):
# 自注意力机制
attn_output = self.self_attn(x, x, x, mask)
x = self.norm1(x + self.dropout(attn_output)) # 残差连接和层归一化
# 前馈网络
ff_output = self.feed_forward(x)
x = self.norm2(x + self.dropout(ff_output)) # 残差连接和层归一化
return x
构建解码器模块
解码器层：包含一个自注意力机制、一个交叉注意力机制和一个前馈网络，每个子层后接残差连接和层归一化。
实例
class DecoderLayer(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout):
super(DecoderLayer, self).__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads) # 自注意力机制
self.cross_attn = MultiHeadAttention(d_model, num_heads) # 交叉注意力机制
self.feed_forward = PositionWiseFeedForward(d_model, d_ff) # 前馈网络
self.norm1 = nn.LayerNorm(d_model) # 层归一化
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout) # Dropout
def forward(self, x, enc_output, src_mask, tgt_mask):
# 自注意力机制
attn_output = self.self_attn(x, x, x, tgt_mask)
x = self.norm1(x + self.dropout(attn_output)) # 残差连接和层归一化
# 交叉注意力机制
attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
x = self.norm2(x + self.dropout(attn_output)) # 残差连接和层归一化
# 前馈网络
ff_output = self.feed_forward(x)
x = self.norm3(x + self.dropout(ff_output)) # 残差连接和层归一化
return x
def __init__(self, d_model, num_heads, d_ff, dropout):
super(DecoderLayer, self).__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads) # 自注意力机制
self.cross_attn = MultiHeadAttention(d_model, num_heads) # 交叉注意力机制
self.feed_forward = PositionWiseFeedForward(d_model, d_ff) # 前馈网络
self.norm1 = nn.LayerNorm(d_model) # 层归一化
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout) # Dropout
def forward(self, x, enc_output, src_mask, tgt_mask):
# 自注意力机制
attn_output = self.self_attn(x, x, x, tgt_mask)
x = self.norm1(x + self.dropout(attn_output)) # 残差连接和层归一化
# 交叉注意力机制
attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
x = self.norm2(x + self.dropout(attn_output)) # 残差连接和层归一化
# 前馈网络
ff_output = self.feed_forward(x)
x = self.norm3(x + self.dropout(ff_output)) # 残差连接和层归一化
return x
构建完整的 Transformer 模型
实例
class Transformer(nn.Module):
def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout):
super(Transformer, self).__init__()
self.encoder_embedding = nn.Embedding(src_vocab_size, d_model) # 编码器词嵌入
self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model) # 解码器词嵌入
self.positional_encoding = PositionalEncoding(d_model, max_seq_length) # 位置编码
# 编码器和解码器层
self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
self.fc = nn.Linear(d_model, tgt_vocab_size) # 最终的全连接层
self.dropout = nn.Dropout(dropout) # Dropout
def generate_mask(self, src, tgt):
# 生成源序列和目标序列的掩码
src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3)
seq_length = tgt.size(1)
nopeak_mask = (1 - torch.triu(torch.ones(1, seq_length, seq_length), diagonal=1)).bool()
tgt_mask = tgt_mask & nopeak_mask
return src_mask, tgt_mask
def forward(self, src, tgt):
# 生成掩码
src_mask, tgt_mask = self.generate_mask(src, tgt)
# 编码器部分
src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
enc_output = src_embedded
for enc_layer in self.encoder_layers:
enc_output = enc_layer(enc_output, src_mask)
# 解码器部分
tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
dec_output = tgt_embedded
for dec_layer in self.decoder_layers:
dec_output = dec_layer(dec_output, enc_output, src_mask, tgt_mask)
# 最终输出
output = self.fc(dec_output)
return output
def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout):
super(Transformer, self).__init__()
self.encoder_embedding = nn.Embedding(src_vocab_size, d_model) # 编码器词嵌入
self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model) # 解码器词嵌入
self.positional_encoding = PositionalEncoding(d_model, max_seq_length) # 位置编码
# 编码器和解码器层
self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
self.fc = nn.Linear(d_model, tgt_vocab_size) # 最终的全连接层
self.dropout = nn.Dropout(dropout) # Dropout
def generate_mask(self, src, tgt):
# 生成源序列和目标序列的掩码
src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3)
seq_length = tgt.size(1)
nopeak_mask = (1 - torch.triu(torch.ones(1, seq_length, seq_length), diagonal=1)).bool()
tgt_mask = tgt_mask & nopeak_mask
return src_mask, tgt_mask
def forward(self, src, tgt):
# 生成掩码
src_mask, tgt_mask = self.generate_mask(src, tgt)
# 编码器部分
src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
enc_output = src_embedded
for enc_layer in self.encoder_layers:
enc_output = enc_layer(enc_output, src_mask)
# 解码器部分
tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
dec_output = tgt_embedded
for dec_layer in self.decoder_layers:
dec_output = dec_layer(dec_output, enc_output, src_mask, tgt_mask)
# 最终输出
output = self.fc(dec_output)
return output
说明：
Transformer 模型：包含编码器和解码器部分，每个部分由多个层堆叠而成。
掩码生成：用于屏蔽无效位置和未来信息。
前向传播：依次通过编码器和解码器，最后通过全连接层输出。
训练 PyTorch Transformer 模型
使用随机数据训练模型，计算损失并更新参数。
实例
# 超参数
src_vocab_size = 5000 # 源词汇表大小
tgt_vocab_size = 5000 # 目标词汇表大小
d_model = 512 # 模型维度
num_heads = 8 # 注意力头数量
num_layers = 6 # 编码器和解码器层数
d_ff = 2048 # 前馈网络内层维度
max_seq_length = 100 # 最大序列长度
dropout = 0.1 # Dropout 概率
# 初始化模型
transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout)
# 生成随机数据
src_data = torch.randint(1, src_vocab_size, (64, max_seq_length)) # 源序列
tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length)) # 目标序列
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss(ignore_index=0) # 忽略填充部分的损失
optimizer = optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)
# 训练循环
transformer.train()
for epoch in range(100):
optimizer.zero_grad()
output = transformer(src_data, tgt_data[:, :-1]) # 目标序列右移
loss = criterion(output.contiguous().view(-1, tgt_vocab_size), tgt_data[:, 1:].contiguous().view(-1))
loss.backward()
optimizer.step()
print(f"Epoch: {epoch+1}, Loss: {loss.item()}")
src_vocab_size = 5000 # 源词汇表大小
tgt_vocab_size = 5000 # 目标词汇表大小
d_model = 512 # 模型维度
num_heads = 8 # 注意力头数量
num_layers = 6 # 编码器和解码器层数
d_ff = 2048 # 前馈网络内层维度
max_seq_length = 100 # 最大序列长度
dropout = 0.1 # Dropout 概率
# 初始化模型
transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_length, dropout)
# 生成随机数据
src_data = torch.randint(1, src_vocab_size, (64, max_seq_length)) # 源序列
tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length)) # 目标序列
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss(ignore_index=0) # 忽略填充部分的损失
optimizer = optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)
# 训练循环
transformer.train()
for epoch in range(100):
optimizer.zero_grad()
output = transformer(src_data, tgt_data[:, :-1]) # 目标序列右移
loss = criterion(output.contiguous().view(-1, tgt_vocab_size), tgt_data[:, 1:].contiguous().view(-1))
loss.backward()
optimizer.step()
print(f"Epoch: {epoch+1}, Loss: {loss.item()}")
模型评估
评估过程：在验证数据上计算损失，评估模型性能。
实例
transformer.eval()
# 生成验证数据
val_src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))
val_tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))
with torch.no_grad():
val_output = transformer(val_src_data, val_tgt_data[:, :-1])
val_loss = criterion(val_output.contiguous().view(-1, tgt_vocab_size), val_tgt_data[:, 1:].contiguous().view(-1))
print(f"Validation Loss: {val_loss.item()}")
# 生成验证数据
val_src_data = torch.randint(1, src_vocab_size, (64, max_seq_length))
val_tgt_data = torch.randint(1, tgt_vocab_size, (64, max_seq_length))
with torch.no_grad():
val_output = transformer(val_src_data, val_tgt_data[:, :-1])
val_loss = criterion(val_output.contiguous().view(-1, tgt_vocab_size), val_tgt_data[:, 1:].contiguous().view(-1))
print(f"Validation Loss: {val_loss.item()}")
点我分享笔记