手撕正弦-余弦位置编码（Sinusoidal Positional Encoding）

news2024/11/26 10:45:55

请添加图片描述

改写后的代码：

import torch
import math
import torch.nn as nn

class PositionalEncoder(nn.Module):
    def __init__(self, d_model, max_seq_len=80):
        super().__init__()
        self.d_model = d_model
        
        # 根据 pos 和 i 创建一个常量 PE 矩阵
        pe = torch.zeros(max_seq_len, d_model)
        for pos in range(max_seq_len):
            for i in range(0, d_model, 2):
                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i) / d_model)))
                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1)) / d_model)))

        # 增加 batch 维度
        pe = pe.unsqueeze(0)
        
        # 将 pe 注册为 buffer，不作为模型的参数，但会在模型中使用
        self.register_buffer('pe', pe)

    def forward(self, x):
        # 使得词嵌入表示相对大一些
        x = x * math.sqrt(self.d_model)
        
        # 获取输入序列的长度
        seq_len = x.size(1)
        
        # 增加位置编码到词嵌入表示中，不需要梯度
        x = x + self.pe[:, :seq_len].cuda()
        
        return x

代码解读：

初始化位置编码矩阵：
- 使用 torch.zeros(max_seq_len, d_model) 初始化了一个大小为 (max_seq_len, d_model) 的零矩阵，之后通过两个循环来填充位置编码的值：
  - 偶数索引维度上用正弦函数生成值。
  - 奇数索引维度上用余弦函数生成值。
unsqueeze(0)：
- 给 pe 增加了一个 batch 维度，以便在输入的 batch 中应用相同的编码矩阵。
register_buffer：
- self.register_buffer('pe', pe)：将 pe 注册为模型的 buffer，它不会在训练过程中更新（即不参与梯度计算），但会作为模型的一部分保存在模型的状态字典中。
前向传播函数 forward：
- 首先通过 x * math.sqrt(self.d_model) 放大输入的词嵌入，这与 Transformer 中的设计一致，用于在数值上平衡词嵌入和位置编码的量级。
- 然后将生成的 pe（位置编码）添加到词嵌入 x 上，self.pe[:, :seq_len].cuda() 取出与输入序列长度相同的部分，并传到 GPU 上（如果在 GPU 上训练）。
注意：这里移除了 Variable，因为现代 PyTorch 不再需要用 Variable 包装张量，直接使用张量即可。由于 pe 不需要计算梯度，在 register_buffer 注册时就默认不参与梯度计算。

如何使用：

你可以将这个类实例化并在模型中使用。例如：

pos_encoder = PositionalEncoder(d_model=512, max_seq_len=100)
x = torch.randn(32, 50, 512)  # 一个 batch 的输入，大小为 (batch_size, seq_len, d_model)
output = pos_encoder(x)  # 加入位置编码

这样就可以为输入序列的每个 token 添加位置编码，并准备好用于 Transformer 模型的输入了。如果你有其他问题，欢迎继续提问！

位置编码的作用：

通过使用正弦和余弦函数，Transformer 的位置编码可以在不同维度上引入不同的周期变化，这样模型能够感知序列中每个 token 的位置以及相对距离。这种编码具有两个优点：

周期性：正弦和余弦的周期性可以很好地表示序列中不同位置之间的关系。
连续性：相邻位置的编码向量非常接近，能够捕捉到相邻 token 之间的顺序信息。

可视化

import numpy as np
import matplotlib.pyplot as plt

# 设置参数
pos = 10  # 位置
d = 512   # 隐藏层维度
i_vals = np.arange(0, d // 2)  # 取一半的维度进行计算

# 计算PE(pos, 2i) 和 PE(pos, 2i+1)
pe_sin = np.sin(pos / (10000 ** (2 * i_vals / d)))
pe_cos = np.cos(pos / (10000 ** (2 * i_vals / d)))

# 画图
plt.figure(figsize=(10, 6))
plt.plot(2 * i_vals, pe_sin, label='sin(pos / 10000^(2i/d))', color='blue')
plt.plot(2 * i_vals + 1, pe_cos, label='cos(pos / 10000^(2i/d))', color='orange')

plt.title(f"Positional Encoding for pos = {pos} and d = {d}")
plt.xlabel('Dimension Index')
plt.ylabel('Value')
plt.legend()
plt.grid(True)
plt.show()

在这里插入图片描述

from mpl_toolkits.mplot3d import Axes3D

# 设置参数
positions = [0, 5, 10, 15, 20]  # 多个位置
d = 512   # 隐藏层维度
i_vals = np.arange(0, d // 2)  # 取一半的维度进行计算

# 创建三维数组保存不同位置的编码
pe_values = np.zeros((len(positions), d))

for idx, pos in enumerate(positions):
    pe_sin = np.sin(pos / (10000 ** (2 * i_vals / d)))
    pe_cos = np.cos(pos / (10000 ** (2 * i_vals / d)))
    pe_values[idx, 2 * i_vals] = pe_sin
    pe_values[idx, 2 * i_vals + 1] = pe_cos

# 画三维图
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 为每个位置画出曲线
for idx, pos in enumerate(positions):
    ax.plot(np.arange(d), [pos]*d, pe_values[idx], label=f'pos={pos}')

# 设置标签
ax.set_xlabel('Dimension Index')
ax.set_ylabel('Position')
ax.set_zlabel('PE Value')
ax.set_title('Positional Encoding for Different Positions')
ax.legend()

plt.show()