File size: 7,497 Bytes
f48983a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | import torch
import math
from torch import nn
class RotaryPositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
"""
初始化 RotaryPositionalEncoding。
Args:
d_model (int): 特征维度。
max_len (int): 支持的最大序列长度。
"""
super(RotaryPositionalEncoding, self).__init__()
# 确保特征维度为偶数
assert d_model % 2 == 0, "d_model must be even for RotaryPositionalEncoding."
# 创建旋转位置编码矩阵
position = torch.arange(0, max_len).float().unsqueeze(1) # [max_len, 1]
dim = torch.arange(0, d_model // 2).float() # [d_model // 2]
div_term = torch.exp(dim * -(math.log(10000.0) / (d_model // 2))) # [d_model // 2]
# 计算正弦和余弦部分
angle = position * div_term # [max_len, d_model // 2]
sin_part = torch.sin(angle) # 正弦部分
cos_part = torch.cos(angle) # 余弦部分
# 将 sin 和 cos 部分堆叠
pe = torch.cat([sin_part, cos_part], dim=-1) # [max_len, d_model]
pe = pe.unsqueeze(0).unsqueeze(0) # [1, 1, max_len, d_model]
self.register_buffer('pe', pe) # 注册为非参数张量
def forward(self, x, offset=0):
"""
前向传播。
Args:
x (Tensor): 输入张量，形状为 [batch_size, seq_len, d_model]。
offset (int): 位置偏移量，默认为 0。
Returns:
Tensor: 应用旋转位置编码的张量。
"""
# 获取位置编码
seq_len = x.size(1)
pe = self.pe[0, :, offset:offset + seq_len, :] # [1, seq_len, d_model]
# 将输入张量拆分为两部分：前半部分与后半部分
x1, x2 = x[..., :x.size(-1)//2], x[..., x.size(-1)//2:] # [batch_size, seq_len, d_model//2]
# 应用旋转操作
x_rotated = torch.cat([
x1 * pe[..., :x.size(-1)//2] - x2 * pe[..., x.size(-1)//2:],
x1 * pe[..., x.size(-1)//2:] + x2 * pe[..., :x.size(-1)//2]
], dim=-1) # [batch_size, seq_len, d_model]
return x_rotated
class ReRoPE:
def __init__(self, dim: int):
"""
初始化 ReRoPE 编码器。
Args:
dim (int): 特征向量的维度（必须为偶数）。
"""
assert dim % 2 == 0, "Dimension must be even for ReRoPE."
self.dim = dim
self.theta = self._compute_base_theta(dim)
@staticmethod
def _compute_base_theta(dim: int):
"""
计算基本的 θ 值，用于旋转位置编码。
Args:
dim (int): 特征向量的维度。
Returns:
torch.Tensor: θ 值的张量。
"""
theta = torch.tensor([10000 ** (-2 * (i // 2) / dim) for i in range(dim)])
return theta
def forward(self, pos: torch.Tensor):
"""
计算给定位置的 ReRoPE 编码。
Args:
pos (torch.Tensor): 位置索引的张量，形状为 [seq_len] 或 [batch_size, seq_len]。
Returns:
torch.Tensor: ReRoPE 编码，形状为 [seq_len, dim] 或 [batch_size, seq_len, dim]。
"""
seq_len = pos.size(-1)
# 获取正弦和余弦部分
angles = pos.unsqueeze(-1) * self.theta
sinusoidal_embedding = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
return sinusoidal_embedding
@staticmethod
def apply_rotary_embedding(query, key, sincos):
"""
应用旋转位置编码到查询和键。
Args:
query (torch.Tensor): 查询向量，形状为 [batch_size, seq_len, dim].
key (torch.Tensor): 键向量，形状为 [batch_size, seq_len, dim].
sincos (torch.Tensor): ReRoPE 编码，形状为 [seq_len, dim] 或 [batch_size, seq_len, dim].
Returns:
Tuple[torch.Tensor, torch.Tensor]: 编码后的查询和键。
"""
sin, cos = sincos[..., :query.size(-1)], sincos[..., query.size(-1):]
query_rotated = query * cos + torch.roll(query, shifts=1, dims=-1) * sin
key_rotated = key * cos + torch.roll(key, shifts=1, dims=-1) * sin
return query_rotated, key_rotated
class LearnablePositionalEmbedding(nn.Module):
def __init__(self, d_model, max_len=5000):
super(LearnablePositionalEmbedding, self).__init__()
# Compute the positional encodings once in log space.
self.pe = nn.Parameter(torch.zeros(
1, 1, max_len, d_model), requires_grad=True)
pe = torch.zeros(max_len, d_model).float()
position = torch.arange(0, max_len).float().unsqueeze(1)
div_term = (torch.arange(0, d_model, 2).float()
* -(math.log(10000.0) / d_model)).exp()
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0).unsqueeze(0)
self.pe.data.copy_(pe.float())
del pe
def forward(self, x, offset=0):
return self.pe[0, :, offset:offset+x.size(1), :]
class SinusoidalPositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
"""
初始化 SinusoidalPositionalEncoding。
Args:
d_model (int): 特征维度。
max_len (int): 支持的最大序列长度。
"""
super(SinusoidalPositionalEncoding, self).__init__()
self.d_model = d_model
# 创建一个固定的位置编码矩阵
pe = self._build_encoding(0, max_len, device=torch.device("cpu"))
# 将矩阵增加批量维度，方便在 forward 中直接使用
pe = pe.unsqueeze(0).unsqueeze(0) # [1, 1, max_len, d_model]
self.register_buffer('pe', pe) # 注册为非参数张量
def _build_encoding(self, start, end, device):
position = torch.arange(start, end, device=device, dtype=torch.float32).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, self.d_model, 2, device=device, dtype=torch.float32)
* -(math.log(10000.0) / self.d_model)
)
pe = torch.zeros(end - start, self.d_model, device=device, dtype=torch.float32)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
return pe
def forward(self, x, offset=0):
"""
前向传播。
Args:
x (Tensor): 输入张量，形状为 [batch_size, seq_len, d_model]
offset (int): 位置偏移量，默认为 0。
Returns:
Tensor: 添加了位置编码的张量。
"""
end = offset + x.size(1)
if end <= self.pe.size(2):
encoding = self.pe[0, :, offset:end, :]
else:
# Some TSQA questions exceed the original fixed 5,000 positions.
# Generate the required range on demand without changing the
# checkpoint buffer shape.
encoding = self._build_encoding(offset, end, x.device).unsqueeze(0)
return encoding.to(device=x.device, dtype=x.dtype)
if __name__ == "__main__":
d_model = 64
seq_len = 128
batch_size = 32
pos_encoder = LearnablePositionalEmbedding(d_model=d_model, max_len=5000)
x = torch.randn(batch_size, seq_len, d_model) # 随机输入张量
pos_encoded_x = pos_encoder(x)
print("Shape of Positional Encoded Output:", pos_encoded_x.shape)
|