ENN Ragged Buffer

此 Python 包实现了一个高效的 RaggedBuffer 数据类型，类似于 3D numpy 数组，但允许第二维具有可变序列长度。它主要用于 enn-trainer，并且目前仅支持 numpy 数组方法的一小部分。

用户指南

使用 pip install ragged-buffer 安装此包。该包目前支持三种 RaggedBuffer 变体，分别是 RaggedBufferF32、RaggedBufferI64 和 RaggedBufferBool。

• 创建 RaggedBuffer
• 获取大小
• 转换为 numpy 数组
• 索引
• 添加
• 连接
• 清除

创建 RaggedBuffer

创建 RaggedBuffer 有三种方式

• RaggedBufferF32(features: int) 创建一个具有指定特征数的空 RaggedBuffer。
• RaggedBufferF32.from_flattened(flattened: np.ndarray, lengths: np.ndarray) 从一个平坦的 2D numpy 数组和一个长度为 1D numpy 数组的长度创建一个 RaggedBuffer。
• RaggedBufferF32.from_array 从一个 3D numpy 数组创建一个 RaggedBuffer（序列长度相等）。

创建一个空的缓冲区并推送每一行

import numpy as np
from ragged_buffer import RaggedBufferF32

# Create an empty RaggedBuffer with a feature size of 3
buffer = RaggedBufferF32(3)
# Push sequences with 3, 5, 0, and 1 elements
buffer.push(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32))
buffer.push(np.array([[10, 11, 12], [13, 14, 15], [16, 17, 18], [19, 20, 21], [22, 23, 24]], dtype=np.float32))
buffer.push(np.array([], dtype=np.float32))  # Alternative: `buffer.push_empty()`
buffer.push(np.array([[25, 25, 27]], dtype=np.float32))

从将第一维和第二维合并的平坦 2D numpy 数组和序列长度数组创建一个 RaggedBuffer

import numpy as np
from ragged_buffer import RaggedBufferF32

buffer = RaggedBufferF32.from_flattened(
    np.array([[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, 25, 27]], dtype=np.float32),
    np.array([3, 5, 0, 1], dtype=np.int64))
)

从具有相同长度的所有序列的 3D numpy 数组创建一个 RaggedBuffer

import numpy as np
from ragged_buffer import RaggedBufferF32

buffer = RaggedBufferF32.from_array(np.zeros((4, 5, 3), dtype=np.float32))

获取大小

size0、size1 和 size2 方法分别返回序列数、序列中的元素数和特征数。

import numpy as np
from ragged_buffer import RaggedBufferF32

buffer = RaggedBufferF32.from_flattened(
    np.zeros((9, 64), dtype=np.float32),
    np.array([3, 5, 0, 1], dtype=np.int64))
)

# Get size of the first/batch dimension.
assert buffer.size0() == 10
# Get size of individual sequences.
assert buffer.size1(1) == 5
assert buffer.size1(2) == 0
# Get size of the last/feature dimension.
assert buffer.size2() == 64

转换为 numpy 数组

as_array 将 RaggedBuffer 转换为一个合并第一维和第二维的平坦 2D numpy 数组。

import numpy as np
from ragged_buffer import RaggedBufferI64

buffer = RaggedBufferI64(1)
buffer.push(np.array([[1], [1], [1]], dtype=np.int64))
buffer.push(np.array([[2], [2]], dtype=np.int64))
assert np.all(buffer.as_array(), np.array([[1], [1], [1], [2], [2]], dtype=np.int64))

索引

您可以使用单个整数来索引 RaggedBuffer，返回一个包含单个序列的 RaggedBuffer，或者使用整数 numpy 数组选择/排列多个序列。

import numpy as np
from ragged_buffer import RaggedBufferF32

# Create a new `RaggedBufferF32`
buffer = RaggedBufferF32.from_flattened(
    np.arange(0, 40, dtype=np.float32).reshape(10, 4),
    np.array([3, 5, 0, 1], dtype=np.int64)
)

# Retrieve the first sequence.
assert np.all(
    buffer[0].as_array() ==
    np.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], dtype=np.float32)
)

# Get a RaggedBatch with 2 randomly selected sequences.
buffer[np.random.permutation(4)[:2]]

添加

如果两个 RaggedBuffer 具有相同的序列数量、序列长度和特征，您可以使用 + 运算符将它们相加。您还可以将所有序列长度为 1 的 RaggedBuffer 添加到具有可变长度序列的 RaggedBuffer 中，沿着每个序列进行广播。

import numpy as np
from ragged_buffer import RaggedBufferF32

# Create ragged buffer with dimensions (3, [1, 3, 2], 1)
rb3 = RaggedBufferI64(1)
rb3.push(np.array([[0]], dtype=np.int64))
rb3.push(np.array([[0], [1], [2]], dtype=np.int64))
rb3.push(np.array([[0], [5]], dtype=np.int64))

# Create ragged buffer with dimensions (3, [1, 1, 1], 1)
rb4 = RaggedBufferI64.from_array(np.array([0, 3, 10], dtype=np.int64).reshape(3, 1, 1))

# Add rb3 and rb4, broadcasting along the sequence dimension.
rb5 = rb3 + rb4
assert np.all(
    rb5.as_array() == np.array([[0], [3], [4], [5], [10], [15]], dtype=np.int64)
)

连接

可以使用 extend 方法通过将其附加到另一个 RaggedBuffer 来修改 RaggedBuffer。

import numpy as np
from ragged_buffer import RaggedBufferF32


rb1 = RaggedBufferF32.from_array(np.zeros((4, 5, 3), dtype=np.float32))
rb2 = RaggedBufferF32.from_array(np.zeros((2, 5, 3), dtype=np.float32))
rb1.extend(r2)
assert rb1.size0() == 6

清除

clear 方法会从 RaggedBuffer 中移除所有元素，而不释放底层内存。

import numpy as np
from ragged_buffer import RaggedBufferF32

rb = RaggedBufferF32.from_array(np.zeros((4, 5, 3), dtype=np.float32))
rb.clear()
assert rb.size0() == 0

许可证

ENN Ragged Buffer dual-licensed under Apache-2.0 and MIT.