1D、2D 和 3D 正弦位置编码（PyTorch 和 TensorFlow）

这是一个实用且易于下载的实现，提供了适用于 PyTorch 和 TensorFlow 的 1D、2D 和 3D 正弦位置编码。

它可以对形状为 (batchsize, x, ch)、(batchsize, x, y, ch) 和 (batchsize, x, y, z, ch) 的张量进行编码，其中位置编码将沿 ch 维度计算。虽然《Attention is All You Need》论文仅支持在一维上进行位置编码，但本实现将其扩展到了二维和三维。

此外，它也适用于形状为 (batchsize, ch, x) 等的张量。更多信息请参阅使用说明。

注意：自版本 6.0.1 起，导入语法已更改。详情请见相关部分。

要安装，只需运行：

pip install positional-encodings[pytorch,tensorflow]

你也可以分别通过以下命令单独安装 PyTorch 和 TensorFlow 的编码模块：

• 仅安装 PyTorch 版本：pip install positional-encodings[pytorch]
• 仅安装 TensorFlow 版本：pip install positional-encodings[tensorflow]

使用说明：

PyTorch

该仓库包含三种主要的位置编码模型：PositionalEncoding{1,2,3}D。此外，还有一个 Summer 类，用于将输入张量与位置编码相加。

import torch
from positional_encodings.torch_encodings import PositionalEncoding1D, PositionalEncoding2D, PositionalEncoding3D, Summer

# 只返回位置编码
p_enc_1d_model = PositionalEncoding1D(10)

# 返回输入与位置编码相加的结果
p_enc_1d_model_sum = Summer(PositionalEncoding1D(10))

x = torch.rand(1,6,10)
penc_no_sum = p_enc_1d_model(x) # penc_no_sum.shape == (1, 6, 10)
penc_sum = p_enc_1d_model_sum(x)
print(penc_no_sum + x == penc_sum) # True

p_enc_2d = PositionalEncoding2D(8)
y = torch.zeros((1,6,2,8))
print(p_enc_2d(y).shape) # (1, 6, 2, 8)

p_enc_3d = PositionalEncoding3D(11)
z = torch.zeros((1,5,6,4,11))
print(p_enc_3d(z).shape) # (1, 5, 6, 4, 11)

对于形状为 (batchsize, ch, x) 或其二维、三维对应形式的张量，在类名前加上 Permute；例如，对于大小为 (batchsize, ch, x) 的 1D 输入，应使用 PositionalEncodingPermute1D 而不是 PositionalEncoding1D。

import torch
from positional_encodings.torch_encodings import PositionalEncodingPermute3D

p_enc_3d = PositionalEncodingPermute3D(11)
z = torch.zeros((1,11,5,6,4))
print(p_enc_3d(z).shape) # (1, 11, 5, 6, 4)

请注意，若需覆盖输出数据类型，可在创建编码时指定：

p_enc_3d = PositionalEncodingPermute3D(11, dtype_override=torch.float64)

这在输入张量为整数类型时特别有用，因为否则输出将始终为零（参见问题 #39）。

TensorFlow Keras

本库同样支持 TensorFlow。只需在所有类名前加上 TF 前缀即可。

import tensorflow as tf
from positional_encodings.tf_encodings import TFPositionalEncoding2D, TFSummer

# 只返回位置编码
p_enc_2d = TFPositionalEncoding2D(170)
y = tf.zeros((1,8,6,2))
print(p_enc_2d(y).shape) # (1, 8, 6, 2)

# 返回输入与位置编码相加的结果
add_p_enc_2d = TFSummer(TFPositionalEncoding2D(170))
y = tf.ones((1,8,6,2))
print(add_p_enc_2d(y) - p_enc_2d(y)) # tf.ones((1,8,6,2))

版本 6.0.1 的变更

在 6.0.1 之前，用户需要同时安装 tensorflow 和 torch 两个包，而这两个包都相当庞大。现在，用户可以单独安装所需的包，但代码也需要相应调整：

如果使用 PyTorch：

from positional_encodings import * -> from positional_encodings.torch_encodings import *

如果使用 TensorFlow：

from positional_encodings import * -> from positional_encodings.tf_encodings import *

公式

插入位置编码的公式如下：

1D：

PE(x,2i) = sin(x/10000^(2i/D))
PE(x,2i+1) = cos(x/10000^(2i/D))

其中：
x 是二维空间中的一个点
i 是 [0, D/2) 内的整数，D 是 ch 维度的大小

2D：

PE(x,y,2i) = sin(x/10000^(4i/D))
PE(x,y,2i+1) = cos(x/10000^(4i/D))
PE(x,y,2j+D/2) = sin(y/10000^(4j/D))
PE(x,y,2j+1+D/2) = cos(y/10000^(4j/D))

其中：
(x,y) 是二维空间中的一个点
i,j 是 [0, D/4) 内的整数，D 是 ch 维度的大小

3D：

PE(x,y,z,2i) = sin(x/10000^(6i/D))
PE(x,y,z,2i+1) = cos(x/10000^(6i/D))
PE(x,y,z,2j+D/3) = sin(y/10000^(6j/D))
PE(x,y,z,2j+1+D/3) = cos(y/10000^(6j/D))
PE(x,y,z,2k+2D/3) = sin(z/10000^(6k/D))
PE(x,y,z,2k+1+2D/3) = cos(z/10000^(6k/D))

其中：
(x,y,z) 是三维空间中的一个点
i,j,k 是 [0, D/6) 内的整数，D 是 ch 维度的大小

3D 公式只是对 这篇 论文中使用的 2D 阐述的自然扩展。

即使输入的维度不能被 2（1D）、4（2D）或 6（3D）整除，也不用担心；必要的填充会自动处理好。

感谢

感谢 这个 仓库为本方法提供了灵感。

引用文献

1D：

@inproceedings{vaswani2017attention,
  title={Attention is all you need},
  author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, Łukasz and Polosukhin, Illia},
  booktitle={Advances in neural information processing systems},
  pages={5998--6008},
  year={2017}
}

2D：

@misc{wang2019translating,
    title={Translating Math Formula Images to LaTeX Sequences Using Deep Neural Networks with Sequence-level Training},
    author={Zelun Wang and Jyh-Charn Liu},
    year={2019},
    eprint={1908.11415},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}

3D： 即将发布！

multidim-positional-encoding 快速上手指南

本工具提供了 PyTorch 和 TensorFlow 版本的 1D、2D 及 3D 正弦位置编码（Sinusoidal Positional Encoding）实现，扩展了经典 Transformer 中仅支持一维编码的限制，适用于处理图像、视频等多维序列数据。

环境准备

• 系统要求：支持 Python 3.6+ 的操作系统（Linux, macOS, Windows）。
• 前置依赖：
  • 必须安装 PyTorch 或 TensorFlow 其中之一（或两者）。
  • 建议预先配置好对应的深度学习框架环境。

安装步骤

推荐使用国内镜像源（如清华源）以加速下载。您可以根据需求选择安装完整包或仅安装特定框架版本。

选项 1：同时安装 PyTorch 和 TensorFlow 支持

pip install positional-encodings[pytorch,tensorflow] -i https://pypi.tuna.tsinghua.edu.cn/simple

选项 2：仅安装 PyTorch 支持（推荐，体积更小）

pip install positional-encodings[pytorch] -i https://pypi.tuna.tsinghua.edu.cn/simple

选项 3：仅安装 TensorFlow 支持

pip install positional-encodings[tensorflow] -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：自版本 6.0.1 起，导入路径已变更。请根据使用的框架从 positional_encodings.torch_encodings 或 positional_encodings.tf_encodings 模块导入，不再直接从根目录导入。

基本使用

PyTorch 使用示例

支持 (batchsize, x, ch)、(batchsize, x, y, ch) 和 (batchsize, x, y, z, ch) 格式的张量。若通道维在中间（如 (batch, ch, x)），请使用带 Permute 的类名。

import torch
from positional_encodings.torch_encodings import PositionalEncoding1D, PositionalEncoding2D, Summer

# 1. 仅获取位置编码
p_enc_1d_model = PositionalEncoding1D(10)
x = torch.rand(1, 6, 10)
penc_no_sum = p_enc_1d_model(x) 
# penc_no_sum.shape == (1, 6, 10)

# 2. 获取输入与位置编码相加的结果 (常用方式)
p_enc_1d_model_sum = Summer(PositionalEncoding1D(10))
penc_sum = p_enc_1d_model_sum(x)
# 验证：penc_sum 等于 x + 位置编码
print(penc_no_sum + x == penc_sum) # True

# 3. 2D 和 3D 示例
p_enc_2d = PositionalEncoding2D(8)
y = torch.zeros((1, 6, 2, 8))
print(p_enc_2d(y).shape) # 输出：(1, 6, 2, 8)

p_enc_3d = PositionalEncoding3D(11)
z = torch.zeros((1, 5, 6, 4, 11))
print(p_enc_3d(z).shape) # 输出：(1, 5, 6, 4, 11)

# 4. 处理通道维在前的情况 (Permute 模式)
# 适用于形状如 (batch, ch, x, y, z) 的张量
from positional_encodings.torch_encodings import PositionalEncodingPermute3D

p_enc_3d_perm = PositionalEncodingPermute3D(11)
z_perm = torch.zeros((1, 11, 5, 6, 4))
print(p_enc_3d_perm(z_perm).shape) # 输出：(1, 11, 5, 6, 4)

# 5. 指定输出数据类型 (避免整数输入导致全零)
p_enc_dtype = PositionalEncodingPermute3D(11, dtype_override=torch.float64)

TensorFlow (Keras) 使用示例

在 TensorFlow 中，所有类名需前缀 TF。

import tensorflow as tf
from positional_encodings.tf_encodings import TFPositionalEncoding2D, TFSummer

# 1. 仅获取位置编码
p_enc_2d = TFPositionalEncoding2D(170)
y = tf.zeros((1, 8, 6, 2))
print(p_enc_2d(y).shape) # 输出：(1, 8, 6, 2)

# 2. 获取输入与位置编码相加的结果
add_p_enc_2d = TFSummer(TFPositionalEncoding2D(170))
y_ones = tf.ones((1, 8, 6, 2))
# 验证：相加后的结果减去纯位置编码应等于原始输入 (全 1)
print(add_p_enc_2d(y_ones) - p_enc_2d(y_ones)) # 输出：tf.ones((1, 8, 6, 2))