注意
转到末尾 下载完整的示例代码。
从流构建 tensordicts¶
作者: Vincent Moens
在许多现实世界的应用中,数据是连续且以不同频率生成的。
例如,来自物联网设备的传感器读数、金融交易或社交媒体更新都可能产生需要实时处理和分析的数据流。
当处理此类数据流时,通常需要将传入的数据“分桶”成离散的块,以便进行高效的处理和分析。然而,当处理具有不同频率或格式的数据流时,这可能具有挑战性。
在本教程中,我们将探讨如何使用 TensorDict 构建和操作数据流。我们将学习如何创建张量的惰性堆叠、处理异步数据流以及为了高效存储和处理而对数据进行密集化。
在本教程中,您将学习: - 如何读取数据流并将它们以固定的时间间隔写入 tensordict 中; - 如何构建将具有异构形状的内容堆叠在一起的 TensorDict; - 如果需要,如何使用 nested_tensor
在单个存储中密集化这些张量。
将异构 tensordicts 堆叠在一起¶
在许多现实场景中,数据以具有不同定义频率的流的形式出现。
本教程的目标是将即将到来的数据“分桶”,以便可以以给定的较慢频率读取和处理它。这种情况下的挑战在于,数据可能无法以规则的“矩形”格式(即,张量的每个维度都明确定义)表示,但可能出现一种情况,即一个数据桶比另一个数据桶具有更多的元素,在这种情况下,我们无法简单地将它们堆叠在一起。通常,考虑以下情况,即第一个和第二个数据桶如下
import torch
from tensordict import TensorDict
bucket0 = TensorDict(stream0=torch.randn(5), stream1=torch.randn(4))
bucket1 = TensorDict(stream0=torch.randn(4), stream1=torch.randn(5))
原则上,我们无法在内存中连续堆叠这两个 tensordict,因为两个流的形状不同。幸运的是,TensorDict 提供了一种工具来将具有异构张量形状的实例分组在一起:LazyStackedTensorDict
。要创建惰性堆叠,只需调用 lazy_stack()
data = TensorDict.lazy_stack([bucket0, bucket1], dim=0)
print(data)
LazyStackedTensorDict(
fields={
stream0: Tensor(shape=torch.Size([2, -1]), device=cpu, dtype=torch.float32, is_shared=False),
stream1: Tensor(shape=torch.Size([2, -1]), device=cpu, dtype=torch.float32, is_shared=False)},
exclusive_fields={
},
batch_size=torch.Size([2]),
device=None,
is_shared=False,
stack_dim=0)
结果数据只是两个 tensordicts 的表示,就好像它们已沿维度 0 堆叠在一起一样。LazyStackedTensorDict
支持 TensorDictBase
类的大多数常见操作,以下是一些示例
data_select = data.select("stream0")
data_plus_1 = data + 1
data_apply = data.apply(lambda x: x + 1)
此外,对其进行索引将返回我们用于创建堆叠的原始数据
assert data[0] is bucket0
尽管如此,在某些情况下,人们可能希望具有底层数据的连续表示。为此,TensorDictBase
提供了一个 densify()
方法,该方法将堆叠可以堆叠的张量,并尝试将其余部分表示为 nested_tensor
实例
data_cont = data.densify()
异步数据流¶
现在让我们切换到一个更具体的示例,在该示例中,我们创建一个以给定频率流式传输数据的函数(在本例中,只是每次迭代递增 1 的整数)。
为了跨线程传递数据,该函数将使用作为输入接收的队列
import asyncio
from typing import List
async def generate_numbers(frequency: float, queue: asyncio.Queue) -> None:
i = 0
while True:
await asyncio.sleep(1 / frequency)
await queue.put(i)
i += 1
collect_data
函数在给定的时间内从队列中读取数据。一旦 timeout
过去,该函数就会返回
async def collect_data(queue: asyncio.Queue, timeout: float) -> List[int]:
values = []
# We create a nested `collect` async function in order to be able to stop it as
# soon as timeout is passed (see wait_for below).
async def collect():
nonlocal values
while True:
value = await queue.get()
values.append(value)
task = asyncio.create_task(collect())
try:
await asyncio.wait_for(task, timeout=timeout)
except asyncio.TimeoutError:
task.cancel()
return values
wait7hz
函数在给定的时间内从队列中读取数据。
async def wait7hz() -> None:
queue = asyncio.Queue()
generate_task = asyncio.create_task(generate_numbers(7, queue))
collect_data_task = asyncio.create_task(collect_data(queue, timeout=1))
values = await collect_data_task
# The ``generate_task`` has not been terminated
generate_task.cancel()
print(values)
asyncio.run(wait7hz())
from typing import Callable, Dict
[0, 1, 2, 3, 4, 5]
我们现在可以设计一个类,该类继承自 LazyStackedTensorDict
并读取来自不同流的数据,并将它们注册到单独的 tensordicts 中。LazyStackedTensorDict
的一个不错的功能是,它也可以增量构建,这样我们只需通过扩展惰性堆叠来注册传入的新数据,直到我们收集到足够的数据为止。以下是 StreamedTensorDict
类的实现
from tensordict import LazyStackedTensorDict, NestedKey, TensorDictBase
class StreamedTensorDict(LazyStackedTensorDict):
"""A lazy stack class that can be built from a dictionary of streams."""
@classmethod
async def from_streams(
cls,
streams: Dict[NestedKey, Callable],
timeout: float,
batch_size: int,
densify: bool = True,
) -> TensorDictBase:
td = cls(stack_dim=0)
# We construct a queue for each stream
queues = [asyncio.Queue() for _ in range(len(streams))]
tasks = []
for stream, queue in zip(streams.values(), queues):
task = asyncio.create_task(stream(queue))
tasks.append(task)
for _ in range(batch_size):
values_tasks = []
for queue in queues:
values_task = asyncio.create_task(collect_data(queue, timeout))
values_tasks.append(values_task)
values = await asyncio.gather(*values_tasks)
td.append(TensorDict(dict(zip(streams.keys(), values))))
# Cancel the generator tasks
for task in tasks:
task.cancel()
if densify:
return td.densify(layout=torch.strided)
return td
最后,main
函数将组合流式传输函数 stream0
和 stream1
,并将它们传递给 StreamedTensorDict.from_streams
方法,该方法将为每个 timeout=1
秒收集 batch_size
批次的数据
async def main() -> TensorDictBase:
def stream0(queue):
return generate_numbers(frequency=7, queue=queue)
def stream1(queue):
return generate_numbers(frequency=3, queue=queue)
# Running this should take about 10 seconds
return await StreamedTensorDict.from_streams(
{"bucket0": stream0, "bucket1": stream1}, timeout=1, batch_size=10
)
td = asyncio.run(main())
print("TensorDict from stream", td)
TensorDict from stream TensorDict(
fields={
bucket0: Tensor(shape=torch.Size([10, -1]), device=cpu, dtype=torch.int64, is_shared=False),
bucket1: Tensor(shape=torch.Size([10, -1]), device=cpu, dtype=torch.int64, is_shared=False)},
batch_size=torch.Size([10]),
device=None,
is_shared=False)
- 让我们表示来自两个流的数据 - 应该等于 torch.arange(),对于 batch_size * timeout * Hz
<=> 1 * 10 秒 * 3 或 7
print("bucket0 (7Hz, around 70 values)", td["bucket0"].values())
print("bucket1 (3Hz, around 30 values)", td["bucket1"].values())
print("shapes of bucket0 (7Hz, around 70 values)", td["bucket0"]._nested_tensor_size())
bucket0 (7Hz, around 70 values) tensor([ 0, 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])
bucket1 (3Hz, around 30 values) tensor([ 0, 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])
shapes of bucket0 (7Hz, around 70 values) tensor([[6],
[7],
[7],
[7],
[7],
[7],
[7],
[7],
[7],
[7]])
结论¶
在本教程中,我们探索了使用 TensorDict 和异步数据流的基础知识。我们学习了如何创建张量的惰性堆叠、如何使用 asyncio 处理异步数据流,以及如何为了高效存储和处理而对数据进行密集化。
我们还了解了如何使用 TensorDict
和 LazyStackedTensorDict
来简化复杂的数据处理任务,例如对具有不同频率的数据流进行分桶。通过利用 TensorDict 和 asyncio 的强大功能,您可以构建可扩展且高效的数据处理管道,这些管道甚至可以处理最苛刻的现实世界应用程序。
感谢您关注本教程!我们希望您发现它对您有所帮助且内容丰富。
脚本的总运行时间: (0 分 11.021 秒)