bash
复制代码
class PositionalEncoding(nn.Module):
def __init__(self, dim, dropout, max_len=5000):
super(PositionalEncoding, self).__init__()
# 确保每个单词的输入维度为偶数，这样sin和cos能配对
if dim % 2 != 0:
raise ValueError("Cannot use sin/cos positional encoding with "
"odd dim (got dim={:d})".format(dim))
"""
构建位置编码pe
pe公式为：
PE(pos,2i/2i+1) = sin/cos(pos/10000^{2i/d_{model}})
"""
pe = torch.zeros(max_len, dim) # max_len 是解码器生成句子的最长的长度，假设是 10 dim为单词的输入维度
# 将位置序号从一维变为只有一列的二维，方便与div_term进行运算，
# 如将[0, 1, 2, 3, 4]变为：
#[
# [0],
# [1],
# [2],
# [3],
# [4]
#]
position = torch.arange(0, max_len).unsqueeze(1)
# 这里使用a^b = e^(blna)公式，来简化运算
# torch.arange(0, dim, 2, dtype=torch.float)表示从0到dim-1，步长为2的一维张量
# 通过以下公式，我们可以得出全部2i的（pos/10000^2i/dim）方便接下来的pe计算
div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) *
-(math.log(10000.0) / dim)))
# 得出的div_term为从0开始，到dim-1，长度为dim/2，步长为2的一维张量
# 将position与div_term做张量乘法，得到的张量形状为（max_len，dim/2）
# 将结果取sin赋给pe中偶数维度，取cos赋给pe中奇数维度
pe[:, 0::2] = torch.sin(position.float() * div_term)
pe[:, 1::2] = torch.cos(position.float() * div_term)
# 将pe的形状从（max_len，dim）变成（max_len，1，dim），在第二个维度上增加一个大小为1的新维度
# 如从原始 pe 张量形状: (5, 4)
#[
# [a1, b1, c1, d1],
# [a2, b2, c2, d2],
# [a3, b3, c3, d3],
# [a4, b4, c4, d4],
# [a5, b5, c5, d5]
#]
# 转换为：执行 unsqueeze(1) 后的 pe 张量形状: (5, 1, 4)
#[
# [[a1, b1, c1, d1]],
# [[a2, b2, c2, d2]],
# [[a3, b3, c3, d3]],
# [[a4, b4, c4, d4]],
# [[a5, b5, c5, d5]]
#]
pe = pe.unsqueeze(1)
# 将pe张量注册为模块的buffer。在PyTorch中，buffer是模型的一部分，但不包含可学习的参数（即不需要梯度）。
# 这样做是因为位置编码在训练过程中是固定的，不需要更新。
self.register_buffer('pe', pe)
self.drop_out = nn.Dropout(p=dropout)
self.dim = dim
def forward(self, emb, step=None):
# 由于位置编码的加入，在我们做softmax时，可能会进入饱和区，使反向传播十分困难，需要进行进行缩放
emb = emb * math.sqrt(self.dim)
if step is None:
emb = emb + self.pe[:emb.size(0)]
else:
emb = emb + self.pe[step]
emb = self.drop_out(emb)
return emb
