位置编码的公式:
PE(pos,2i)=sin(pos100002i/dmodel)PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)PE(pos,2i)=sin(100002i/dmodelpos)PE(pos,2i+1)=cos(pos100002i/dmodel)PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)PE(pos,2i+1)=cos(100002i/dmodelpos)
pospospos:词在句子中的位置(0,1,2,...0, 1, 2, ...0,1,2,...)。
iii:位置向量的维度索引(0,1,2,...,2550, 1, 2, ... , 2550,1,2,...,255)。
2i2i2i表示偶数维度(第 0, 2, 4 维用正弦sin\sinsin),2i+12i+12i+1表示奇数维度(第 1, 3, 5 维用余弦cos\coscos)。
importtorchimporttorch.nnasnnimportmath# 1. 位置编码 (Positional Encoding)classPositionalEncoding(nn.Module):def__init__(self,d_model,max_len=5000):super(PositionalEncoding,self).__init__()# 创建一个足够长的矩阵来装位置编码pe=torch.zeros(max_len,d_model)position=torch.arange(0,max_len,dtype=torch.float).unsqueeze(1)# 计算 10000^(2i/d_model)div_term=torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))# 偶数维度用 sin,奇数维度用 cos (对应上面提到的表盘原理)pe[:,0::2]=torch.sin(position*div_term)pe[:,1::2]=torch.cos(position*div_term)# 增加一个 batch 维度:(1, max_len, d_model)pe=pe.unsqueeze(0)# 注册为 buffer,这样它就不会被当作模型参数被反向传播更新self.register_buffer('pe',pe)将位置编码直接加到词向量上
Input=WordEmbedding(x)+PositionalEncoding(pos)\text{Input} = \text{WordEmbedding}(x) + \text{PositionalEncoding}(pos)Input=WordEmbedding(x)+PositionalEncoding(pos)
defforward(self,x):""" x 的 shape: (batch_size, seq_len, d_model) 把位置编码直接加到词向量上 """x=x+self.pe[:,:x.size(1),:]returnx# 2. 核心架构:Transformer 编码器模型classSimpleTransformerEncoder(nn.Module):def__init__(self,vocab_size,d_model,nhead,num_layers,dim_feedforward,max_seq_len):""" vocab_size: 词汇表大小 d_model: 词向量维度 (通常也是 QKV 的维度) nhead: 多头注意力的“头”数 num_layers: 编码器 Block 堆叠的层数 dim_feedforward: 前馈神经网络隐藏层的维度 """super(SimpleTransformerEncoder,self).__init__()self.d_model=d_model# A. 将词汇 ID 映射为词向量 (Input Embedding)self.embedding=nn.Embedding(vocab_size,d_model)# B. 加入位置编码self.pos_encoder=PositionalEncoding(d_model,max_seq_len)# C. 定义单层 Encoder Block 的内部结构 (PyTorch 已高度封装)encoder_layer=nn.TransformerEncoderLayer(d_model=d_model,nhead=nhead,dim_feedforward=dim_feedforward,batch_first=True# 设定输入格式为 (batch, seq, feature))# D. 将上面的单层 Block 堆叠 num_layers 次self.transformer_encoder=nn.TransformerEncoder(encoder_layer,num_layers=num_layers)# E. (可选) 假设我们在做一个分类任务,最后接一个全连接层self.fc_out=nn.Linear(d_model,2)# 比如输出二分类# 执行transformerdefforward(self,src,src_key_padding_mask=None):""" src: 输入序列的 Token ID,shape: (batch_size, seq_len) src_key_padding_mask: 用于遮挡填充的 <PAD> token,不让注意力计算它们 """# 1. 生成词向量并乘以根号 d_model (论文中的标准化技巧)src=self.embedding(src)*math.sqrt(self.d_model)# 2. 加上位置编码src=self.pos_encoder(src)# 3. 送入 Transformer 编码器核心# 输出的 memory 包含了句子里每个词深度融合后的上下文特征memory=self.transformer_encoder(src,src_key_padding_mask=src_key_padding_mask)# 4. 如果是分类任务,通常取第一个 token (类似 BERT 的 [CLS]) 或求平均# 这里为了演示,我们对整个序列的时间步(seq_len)维度求平均pooled_output=memory.mean(dim=1)# 5. 输出最终预测结果logits=self.fc_out(pooled_output)returnlogits