FoldedTensor: PyTorch扩展，用于处理深度嵌套的可变长度序列

foldedtensor是一个PyTorch扩展，提供对包含深度嵌套可变大小序列的张量的高效处理。它允许根据序列长度的内部结构来展开/折叠（或展开/折叠）数据维度。当处理可以以不同方式分割的数据时，该库特别有用，并允许您避免选择固定的表示。

安装

可以使用pip安装库

pip install foldedtensor

特性

• 支持任意数量的嵌套维度
• 处理已填充张量时无计算开销
• 基于存储的内部长度动态重新填充（或重新折叠）数据
• 每次张量重新折叠时自动生成和更新掩码
• C++优化代码，用于快速从Python列表加载数据和折叠
• 数据表示的灵活性，便于在需要时切换到不同的布局

示例

在最简单的情况下，foldedtensor可以用于将嵌套Python列表转换为PyTorch张量

from foldedtensor import as_folded_tensor

ft = as_folded_tensor(
    [
        [0, 1, 2],
        [3],
    ],
)
# FoldedTensor([[0, 1, 2],
#               [3, 0, 0]])

您也可以在创建时指定名称和平铺/非平铺维度

import torch
from foldedtensor import as_folded_tensor

# Creating a folded tensor from a nested list
# There are 2 samples, the first with 5 lines, the second with 1 line.
# Each line contain between 1 and 2 words.
ft = as_folded_tensor(
    [
        [[1], [], [], [], [2, 3]],
        [[4, 3]],
    ],
    data_dims=("samples", "words"),
    full_names=("samples", "lines", "words"),
    dtype=torch.long,
)
print(ft)
# FoldedTensor([[1, 2, 3],
#               [4, 3, 0]])

创建后，您可以通过折叠来改变张量的形状

# Refold on the lines and words dims (flatten the samples dim)
print(ft.refold(("lines", "words")))
# FoldedTensor([[1, 0],
#               [0, 0],
#               [0, 0],
#               [0, 0],
#               [2, 3],
#               [4, 3]])

# Refold on the words dim only: flatten everything
print(ft.refold(("words",)))
# FoldedTensor([1, 2, 3, 4, 3])

张量可以进一步与标准PyTorch操作一起使用

# Working with PyTorch operations
embedder = torch.nn.Embedding(10, 16)
embedding = embedder(ft.refold(("words",)))
print(embedding.shape)
# torch.Size([5, 16]) # 5 words total, 16 dims

refolded_embedding = embedding.refold(("samples", "words"))
print(refolded_embedding.shape)
# torch.Size([2, 5, 16]) # 2 samples, 5 words max, 16 dims

基准测试

在此处查看与各种替代方案的比较：基准测试文档。

与替代方案的比较

与其他的锯齿或嵌套张量实现不同，FoldedTensor不会对嵌套数据进行特定结构的约束，也不需要填充所有维度。这在使用可以按数据转换的不同方式排列的数据时，为用户提供更大的灵活性。此外，C++优化确保了高性能，使其非常适合高效处理深层嵌套张量。

以下是与其他常见处理可变长度嵌套序列实现的比较