Transformer代码实现2：手搓词嵌入层和位置编码

为读者更好的阅读体验，请跳转至：https://1oqodyst.html2web.com

一、词嵌入层在 Transformer 架构中的位置

从架构图中可以看出，词嵌入层的主要作用是在一开始就将输入（inputs and outputs）的token 序列转为一个编码的向量，同时还需要加上位置编码。

二、具体代码实现

2.1 导入相关库

首先还是导入相关库

import math import torch import torch.nn as nn from torch.autograd import Variable ​ from parser1 import args

这里需要介绍一下from torch.autograd import Variable，它的作用是让张量支持自动求导。从 PyTorch 0.4 版本开始，Tensor和Variable已经被合并。现在，torch.Tensor本身就集成了自动求导的所有功能。你不再需要使用Variable进行包装，只需在创建Tensor时直接设置requires_grad=True即可。

parser1.py中实现了相关参数配置，这里省略，后面代码中看到的包含args.的参数都是该文件中配置好的。

2.2 实现词嵌入层类

依旧继承于nn.Module

class Embeddings(nn.Module):

__init__函数：

其参数包括：

• d_{model}: 模型的维度（嵌入向量的维度），论文中为512

• vocab: 词表大小

def __init__(self, d_model, vocab): super(Embeddings, self).__init__() # 标准的构造函数写法 self.lut = nn.Embedding(vocab, d_model) self.d_model = d_model

LUT是Look-Up Table的缩写，即查找表。

nn.Embedding是 PyTorch 内置的查找表，其内容是随机初始化的，形状为(vocab, d_model)

forward函数

def forward(self, x): return self.lut(x) * math.sqrt(self.d_model)

在 Transformer 中，嵌入向量随后会与位置编码相加。为了防止嵌入向量的数值过大或过小，导致后续的 Softmax 函数梯度变得极小（从而无法训练），作者建议使用 d_{model} 进行缩放，以保持方差的稳定。

2.3 实现位置编码层类

依旧继承于nn.Module

class PositionalEncoding(nn.Module):

__init__函数：

def __init__(self, d_model, dropout, max_len=5000): super(PositionalEncoding, self).__init__() # 标准的构造函数写法 self.dropout = nn.Dropout(p=dropout) ​ # 初始化位置编码矩阵 # pe形状: (max_len, d_model)，每一行代表一个位置的编码 pe = torch.zeros(max_len, d_model, device=args.device) # position: 位置索引向量，形状 (max_len, 1) # [[0.], [1.], [2.], ..., [max_len-1.]] position = torch.arange(0., max_len, device=args.device).unsqueeze(1) # div_term: 除数项，用于计算不同维度的频率 # 公式: 1 / 10000^(2i/d_model) = exp(-log(10000) * 2i / d_model) # 形状: (d_model/2,) # 例如 d_model=512 时: [1.0, 0.95, 0.90, ..., 0.0001] # 不同维度使用不同的频率，低维度频率高，高维度频率低 div_term = torch.exp(torch.arange(0., d_model, 2, device=args.device) * -(math.log(10000.0) / d_model)) # 偶数维度使用sin函数 # pe[:, 0::2] 表示取所有行的第0, 2, 4, ...列 # position * div_term 广播后形状为 (max_len, d_model/2) pe[:, 0::2] = torch.sin(position * div_term) # 奇数维度使用cos函数 # pe[:, 1::2] 表示取所有行的第1, 3, 5, ...列 pe[:, 1::2] = torch.cos(position * div_term) # 在第0维增加一个batch维度 # pe形状从 (max_len, d_model) 变为 (1, max_len, d_model) # 这样便于后续与输入x (batch_size, seq_len, d_model) 进行广播相加 pe = pe.unsqueeze(0) # register_buffer: 将pe注册为模型的buffer # buffer与parameter的区别: # - parameter: 会被优化器更新（可学习参数） # - buffer: 不会更新，但会随模型移动到GPU/CPU，也会保存到checkpoint # 位置编码是固定的，不需要学习，所以用buffer self.register_buffer('pe', pe)

频率项计算公式：

$$
div\_term = \exp\left(-\log(10000) \cdot \frac{2i}{d_{model}}\right) = \frac{1}{10000^{\frac{2i}{d_{model}}}}
$$

其中代码torch.arange(0., d_model, 2, device=args.device)主要实现 2i 的计算。

假设d_model = 10，这行代码会生成如下张量：[0.,2.,4.,6.,8.]

位置编码计算公式：

$$
PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{\frac{2i}{d_{model}}}}\right) \\ PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{\frac{2i}{d_{model}}}}\right)
$$

forward函数

参数​x:​ 词嵌入向量，形状(batch_size, seq_len, d_model)

最后返回添加位置编码后的向量，形状不变。

def forward(self, x): x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False) return self.dropout(x)

self.pe[:, :x.size(1)]: 取前seq_len个位置编码形状从 (1, max_len, d_model) 变为 (1, seq_len, d_model)

Variable(..., requires_grad=False): 包装为不需要梯度的变量（位置编码是固定的，不需要计算梯度）

x + pe: 广播相加，pe会自动扩展到batch_size

dropout: 随机置零部分元素，防止过拟合

2.4 完整代码如下

import math import torch import torch.nn as nn from torch.autograd import Variable ​ from parser1 import args ​ ​ class Embeddings(nn.Module): def __init__(self, d_model, vocab): super(Embeddings, self).__init__() self.lut = nn.Embedding(vocab, d_model) self.d_model = d_model ​ def forward(self, x): return self.lut(x) * math.sqrt(self.d_model) ​ ​ class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout, max_len=5000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(p=dropout) pe = torch.zeros(max_len, d_model, device=args.device) position = torch.arange(0., max_len, device=args.device).unsqueeze(1) div_term = torch.exp(torch.arange(0., d_model, 2, device=args.device) * -(math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) self.register_buffer('pe', pe) ​ def forward(self, x): x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False) return self.dropout(x) ​