import torch
import torch.nn as nn
import torch.nn.functional as F
import math
# 定义位置编码类,用于为Transformer模型中的输入序列添加位置信息
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super(PositionalEncoding, self).__init__()
# 初始化一个全零张量,用于存储位置编码
pe = torch.zeros(max_len, d_model)
# 创建一个从0到max_len-1的位置序列
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# 计算用于位置编码的缩放因子
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
# 使用正弦函数计算偶数位置的位置编码
pe[:, 0::2] = torch.sin(position * div_term)
# 使用余弦函数计算奇数位置的位置编码
pe[:, 1::2] = torch.cos(position * div_term)
# 调整张量形状以匹配Transformer模型的输入形状
pe = pe.unsqueeze(0).transpose(0, 1)
# 将位置编码张量注册为缓冲区,这样它就不会被模型的参数优化器更新
self.register_buffer('pe', pe)
def forward(self, x):
# 将位置编码添加到输入张量上
x = x + self.pe[:x.size(0), :]
return x
# 定义多头自注意力机制类
class MultiHeadAttention(nn.Module):
# 初始化函数,设置模型维度和头数等参数
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
self.num_heads = num_heads
self.d_k = d_model // num_heads # 每个头的维度
# 定义用于线性变换的层,生成查询、键和值张量
self.q_linear = nn.Linear(d_model, d_model)
self.k_linear = nn.Linear(d_model, d_model)
self.v_linear = nn.Linear(d_model, d_model)
self.out_linear = nn.Linear(d_model, d_model)
# 前向传播函数,实现多头自注意力机制的计算过程
def forward(self, query, key, value, mask=None):
# 获取批次大小
batch_size = query.size(0)
# 通过线性变换生成查询、键和值张量,并进行形状调整以适应多头处理
Q = self.q_linear(query).view(batch_size, -1, self.num_heads, self.d_k).permute(0, 2, 1, 3)
K = self.k_linear(key).view(batch_size, -1, self.num_heads, self.d_k).permute(0, 2, 1, 3)
V = self.v_linear(value).view(batch_size, -1, self.num_heads, self.d_k).permute(0, 2, 1, 3)
# 计算注意力分数,并应用缩放因子和可选的掩码
scores = torch.matmul(Q, K.permute(0, 1, 3, 2)) / math.sqrt(self.d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9) # 将掩码为0的位置的分数设置为极小值
# 通过softmax函数计算注意力权重,并与值张量相乘得到上下文张量
attn_weights = F.softmax(scores, dim=-1)
context = torch.matmul(attn_weights, V).permute(0, 2, 1, 3).contiguous().view(batch_size, -1, self.d_k * self.num_heads)
# 通过线性变换输出最终的上下文表示
output = self.out_linear(context)
return output
# 定义Transformer块类,包含自注意力机制和前馈神经网络等组件
class TransformerBlock(nn.Module):
# 初始化函数,设置模型维度、头数、dropout概率和前馈神经网络维度等参数
def __init__(self, d_model, num_heads, dropout=0.1, ff_dim=2048):
super(TransformerBlock, self).__init__()
# 定义自注意力机制层、层归一化层和dropout层等组件
self.self_attn = MultiHeadAttention(d_model, num_heads) # MultiHeadAttention类实例化
self.norm1 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(dropout)
# 定义前馈神经网络,包含两个线性层和一个ReLU激活函数
self.ffn = nn.Sequential(
nn.Linear(d_model, ff_dim),
nn.ReLU(),
nn.Linear(ff_dim, d_model)
)
self.norm2 = nn.LayerNorm(d_model)
self.dropout2 = nn.Dropout(dropout)
# 前向传播函数,实现Transformer块的计算过程
def forward(self, src, src_mask=None):
# 通过自注意力机制层处理输入张量,并添加残差连接和层归一化
src2 = self.self_attn(src, src, src, mask=src_mask) # 多头注意力
src = src + self.dropout1(src2)
src = self.norm1(src)
# 通过前馈神经网络处理张量,并再次添加残差连接和层归一化
src2 = self.ffn(src)
# 对 src2 的输出应用了一个Dropout操作。Dropout是一种常用的正则化手段,它可以随机“关闭”一部分神经元,防止模型过度依赖某一部分特征,以此增强模型的泛化能力。
src = src + self.dropout2(src2) # 一个带Dropout的残差连接
src = self.norm2(src)
return src
# 定义完整的Transformer模型类,包含多个Transformer块和位置编码等组件
class TransformerModel(nn.Module):
# 初始化函数,设置输入维度、模型维度、头数、层数、dropout概率和前馈神经网络维度等参数
def __init__(self, input_dim, d_model, num_heads, num_layers, dropout=0.1, ff_dim=2048):
super(TransformerModel, self).__init__()
# 定义位置编码层、多个Transformer块和全连接输出层等组件
self.positional_encoding = PositionalEncoding(d_model) # 位置编码实例化
# nn.ModuleList存储一组可学习参数的神经网络层或者其他任何继承自 nn.Module 类的模块
self.encoder_layers = nn.ModuleList([TransformerBlock(d_model, num_heads, dropout, ff_dim) for _ in range(num_layers)])
self.fc_out = nn.Linear(d_model, input_dim)
# 前向传播函数,实现Transformer模型的完整计算过程
def forward(self, src, src_mask=None):
# 为输入张量添加位置编码,并通过多个Transformer块进行处理
src = self.positional_encoding(src)
for layer in self.encoder_layers:
src = layer(src, src_mask)
# 通过全连接输出层生成最终的输出张量
output = self.fc_out(src)
return output
# 使用示例
# 创建一个Transformer模型实例并指定相关参数
input_dim = 100 # 输入维度为100维(例如:词嵌入的维度)
d_model = 512 # 模型内部使用的维度为512维
num_heads = 8 # 使用8个头进行多头自注意力机制的计算
num_layers = 6 # 使用6个Transformer块堆叠构建模型
dropout = 0.1 # dropout概率为0.1,用于防止过拟合
ff_dim = 2048 # 前馈神经网络的隐藏层维度为2048维
# 实例化Transformer模型,并传入上述参数进行初始化
model = TransformerModel(input_dim, d_model, num_heads, num_layers, dropout, ff_dim)
print(model)
结构:
该模型首先对输入序列加上位置编码,然后通过6层TransformerBlock进行编码处理,每层包含自注意力机制和前馈神经网络结构,同时配合层归一化和Dropout正则化,最后通过一个全连接层输出100维的特征向量。
-
PositionalEncoding:
- 位置编码层,用于给输入序列的每个位置附加位置信息,以便模型能够捕捉到序列中元素的位置依赖关系。
-
encoder_layers:
- 这是一个
ModuleList
容器,包含了6个TransformerBlock
模块,每个模块对应编码器的一层。
- 这是一个
-
TransformerBlock:
- 每个
TransformerBlock
由以下几部分组成:-
self_attn:
-
MultiHeadAttention
模块,负责对输入序列执行自注意力机制,包括:q_linear
、k_linear
、v_linear
:分别用于查询(Query)、键(Key)和值(Value)的线性变换,输入和输出维度均为512。out_linear
:将多头注意力的输出进一步映射回原始维度的线性层。
-
-
norm1:
LayerNorm
层,对自注意力输出进行层归一化,稳定模型训练并提高性能。
-
dropout1:
Dropout
层,引入随机失活以防止过拟合并增加模型的泛化能力。
-
ffn:
Sequential
容器,包含一个两层的前馈神经网络(FFN):- 第一层是
Linear
层,将输入的512维特征映射到2048维。 - 第二层是
ReLU
激活函数,接着是一层将2048维特征映射回512维的Linear
层。
- 第一层是
-
norm2:
- 另一个
LayerNorm
层,对FFN的输出进行归一化。
- 另一个
-
dropout2:
- 第二个
Dropout
层,作用同上。
- 第二个
-
- 每个
-
fc_out:
- 最后一个
Linear
层(全连接层),将编码器最后一层输出的512维向量映射到100维,作为最终的输出结果。
- 最后一个
TransformerModel(
(positional_encoding): PositionalEncoding()
(encoder_layers): ModuleList(
(0-5): 6 x TransformerBlock(
(self_attn): MultiHeadAttention(
(q_linear): Linear(in_features=512, out_features=512, bias=True)
(k_linear): Linear(in_features=512, out_features=512, bias=True)
(v_linear): Linear(in_features=512, out_features=512, bias=True)
(out_linear): Linear(in_features=512, out_features=512, bias=True)
)
(norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
(dropout1): Dropout(p=0.1, inplace=False)
(ffn): Sequential(
(0): Linear(in_features=512, out_features=2048, bias=True)
(1): ReLU()
(2): Linear(in_features=2048, out_features=512, bias=True)
)
(norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
(dropout2): Dropout(p=0.1, inplace=False)
)
)
(fc_out): Linear(in_features=512, out_features=100, bias=True)
)