PyTorch中是否内置positional encoding？如何指定维度并获取对应编码

PyTorch官方并没有内置专门的位置编码（positional encoding）模块，但你可以快速实现一个满足需求的版本——既能指定编码维度，也能直接获取任意第i个位置对应的编码值。

最常用的是Transformer论文中提出的正弦余弦位置编码，它无需训练，直接通过公式计算生成，完全符合你的需求。以下是具体实现：

import torch
import math

class PositionalEncoding(torch.nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        # 初始化位置编码矩阵
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)
    
    def get_position_encoding(self, idx: int) -> torch.Tensor:
        # 获取第idx个位置的编码（idx从0开始）
        return self.pe[idx].squeeze(0)

# 使用示例
d_model = 512  # 指定编码维度
pe = PositionalEncoding(d_model=d_model)

# 获取第10个位置的编码
pos_10 = pe.get_position_encoding(10)
print(pos_10.shape)  # 输出: torch.Size([512])

关键说明： #

• 指定编码维度：通过d_model参数直接设置编码的维度大小，和你的输入特征维度保持一致即可。
• 获取任意位置编码：调用get_position_encoding(idx)方法，传入位置索引（从0开始）就能得到对应位置的编码张量。
• 无需训练：正弦余弦编码是固定计算的，不需要参与反向传播，适合需要直接获取任意位置编码的场景。

如果你需要可学习的位置编码（编码值随训练更新），也可以简单修改实现：

class LearnablePositionalEncoding(torch.nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        self.pe = torch.nn.Parameter(torch.randn(max_len, 1, d_model))
    
    def get_position_encoding(self, idx: int) -> torch.Tensor:
        return self.pe[idx].squeeze(0)

这种版本的编码值会在训练中优化，同样支持指定维度和获取任意位置的编码。

内容的提问来源于stack exchange，提问作者Bipolo