tensordict 包¶
该 TensorDict
类通过打包一个继承自常规 pytorch 张量特性的类字典对象,简化了在不同模块之间传递多个张量的过程。
|
TensorDictBase 是 TensorDict 的抽象父类,TensorDict 是一个 torch.Tensor 数据容器。 |
|
一个批处理的张量字典。 |
|
一个惰性堆叠的 TensorDict。 |
|
持久化 TensorDict 实现。 |
|
一个具有参数暴露功能的 TensorDictBase 包装器。 |
|
返回 get 默认值的状态。 |
构造函数和处理器¶
该库提供了几种与 numpy 结构化数组、命名元组或 h5 文件等其他数据结构进行交互的方法。该库还暴露了专门用于操作 tensordicts 的函数,如 save
、load
、stack
或 cat
。
|
将tensordicts沿给定维度连接成一个tensordict。 |
|
当一个类型不是张量集合(tensordict 或 tensorclass)时,返回 |
|
将任何对象转换为 TensorDict。 |
|
从合并文件重建 tensordict。 |
|
将字典转换为 TensorDict。 |
|
将 HDF5 文件转换为 TensorDict。 |
|
将模块的参数和缓冲区复制到 tensordict 中。 |
|
为 vmap 的 ensemable 学习/特征期望应用检索多个模块的参数。 |
|
将命名元组转换为 TensorDict。 |
|
将 pytree 转换为 TensorDict 实例。 |
|
将结构化 numpy 数组转换为 TensorDict。 |
|
将元组转换为 TensorDict。 |
|
从键列表和单个值创建 tensordict。 |
|
|
|
当一个类型不是张量集合(tensordict 或 tensorclass)或是一个非张量时,返回 |
|
创建 TensorDicts 的懒惰堆叠。 |
|
从磁盘加载 tensordict。 |
|
从磁盘加载内存映射的 tensordict。 |
|
尝试使 TensorDicts 密集堆叠,并在需要时回退到懒惰堆叠。 |
|
将所有张量写入内存映射的 Tensor 中,并放入新的 tensordict。 |
|
将tensordict保存到磁盘。 |
|
沿给定维度将 tensordicts 堆叠成一个单一的 tensordict。 |
TensorDict 作为上下文管理器¶
在需要执行某个操作然后撤销该操作的情况下,TensorDict
可以用作上下文管理器。这包括临时锁定/解锁一个 tensordict。
>>> data.lock_() # data.set will result in an exception
>>> with data.unlock_():
... data.set("key", value)
>>> assert data.is_locked()
或者,在包含模型参数和缓冲区的 TensorDict 实例的帮助下执行函数调用。
>>> params = TensorDict.from_module(module).clone()
>>> params.zero_()
>>> with params.to_module(module):
... y = module(x)
在第一个例子中,因为我们暂时解锁了 tensordict data,所以可以修改它。在第二个例子中,我们用 params tensordict 实例中包含的参数和缓冲区填充模型,并在该调用完成后重置原始参数。
内存映射张量¶
tensordict 提供了 MemoryMappedTensor
原始类型,允许您以一种方便的方式处理存储在物理内存中的张量。 MemoryMappedTensor
的主要优点是其易于构建(无需处理张量的存储)、可以处理不适合内存的大型连续数据、跨进程的高效(反)序列化以及对存储张量的高效索引。
如果所有工作进程都能访问相同的存储(无论是在多进程还是分布式环境中),传递 MemoryMappedTensor
只需传递一个指向磁盘上文件的引用以及一些用于重建它的额外元数据。对于索引内存映射张量也是如此,只要它们的存储的数据指针与原始指针相同。
索引内存映射张量比从磁盘加载多个独立文件要快得多,并且不需要将数组的全部内容加载到内存中。但是,PyTorch 张量的物理存储不应有任何区别。
>>> my_images = MemoryMappedTensor.empty((1_000_000, 3, 480, 480), dtype=torch.unint8)
>>> mini_batch = my_images[:10] # just reads the first 10 images of the dataset
|
一个内存映射张量。 |
逐点运算¶
Tensordict 支持各种逐点运算,允许您对其中存储的张量执行按元素计算。这些运算类似于对常规 PyTorch 张量执行的运算。
支持的运算¶
当前支持以下逐点运算:
左加和右加(+)
左减和右减(-)
左乘和右乘(*)
左除和右除(/)
左幂(**)
执行逐点运算¶
您可以在两个 Tensordicts 之间,或者在 Tensordict 与张量/标量值之间执行逐点运算。
示例 1:Tensordict-Tensordict 运算¶
>>> import torch
>>> from tensordict import TensorDict
>>> td1 = TensorDict(
... a=torch.randn(3, 4),
... b=torch.zeros(3, 4, 5),
... c=torch.ones(3, 4, 5, 6),
... batch_size=(3, 4),
... )
>>> td2 = TensorDict(
... a=torch.randn(3, 4),
... b=torch.zeros(3, 4, 5),
... c=torch.ones(3, 4, 5, 6),
... batch_size=(3, 4),
... )
>>> result = td1 * td2
在此示例中,\* 运算符按元素应用于 td1 和 td2 中的相应张量。
示例 2:Tensordict-Tensor 运算¶
>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(
... a=torch.randn(3, 4),
... b=torch.zeros(3, 4, 5),
... c=torch.ones(3, 4, 5, 6),
... batch_size=(3, 4),
... )
>>> tensor = torch.randn(4)
>>> result = td * tensor
在此,\* 运算符按元素应用于 td 中的每个张量和提供的张量。该张量被广播以匹配 Tensordict 中每个张量的形状。
示例 3:Tensordict-Scalar 运算¶
>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(
... a=torch.randn(3, 4),
... b=torch.zeros(3, 4, 5),
... c=torch.ones(3, 4, 5, 6),
... batch_size=(3, 4),
... )
>>> scalar = 2.0
>>> result = td * scalar
在这种情况下,\* 运算符按元素应用于 td 中的每个张量和提供的标量。
广播规则¶
当在 Tensordict 与张量/标量之间执行逐点运算时,张量/标量会被广播以匹配 Tensordict 中每个张量的形状:张量在左侧广播以匹配 tensordict 形状,然后单独在右侧广播以匹配张量形状。这遵循 PyTorch 中使用的标准广播规则,如果您将 TensorDict
视为单个张量实例。
例如,如果您有一个形状为 (3, 4)
的张量的 Tensordict,并将其乘以形状为 (4,)
的张量,则该张量将在应用运算之前被广播到形状 (3, 4)。如果 tensordict 包含一个形状为 (3, 4, 5)
的张量,用于乘法的张量将被广播到右侧的 (3, 4, 5)
以进行该乘法。
如果逐点运算在多个 tensordicts 之间执行,并且它们的批次大小不同,那么它们将被广播到一个共同的形状。
逐点运算的效率¶
如果可能,将使用 torch._foreach_<op>
融合内核来加速逐点运算的计算。
处理缺失条目¶
当在两个 Tensordicts 之间执行逐点运算时,它们必须具有相同的键。某些运算,如 add()
,有一个 default
关键字参数,可用于与具有独有条目的 tensordict 进行运算。如果 default=None
(默认值),则两个 Tensordicts 必须具有完全匹配的键集。如果 default="intersection"
,则只考虑相交的键集,其他键将被忽略。在所有其他情况下,default
将用于运算双方的所有缺失条目。
torch.autograd.grad
的梯度计算¶
Tensordict 扩展了 :pyfunc:`torch.autograd.grad`,以便可以直接在(嵌套的)TensorDict
结构上计算梯度。所有张量都被提取为元组,执行底层的 PyTorch 操作,然后结果被打包回具有相同键结构的新 TensorDict
中。这使得区分处理一次处理多个张量的复杂管道变得简单明了。
简单示例¶
>>> from tensordict import TensorDict
>>> import torch
>>> # Build input TensorDict with requires_grad
>>> inputs = TensorDict(a=torch.randn(2, 3, requires_grad=True))
>>> outputs = inputs + 1.0
>>> grads = torch.autograd.grad(outputs, inputs, torch.ones_like(outputs))
>>> print(grads)
TensorDict(
fields={
a: Tensor(shape=torch.Size([2, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([]),
device=None,
is_shared=False)
端到端模块¶
该特性在处理接受 tensordict 并返回另一个 tensordict 的函数式模块时(例如通过 tensordict.nn.TensorDictModule()
)非常有用。
>>> td_in = TensorDict(a=torch.randn(2, 3, requires_grad=True))
>>> td_out = td_module(td_in)
>>> td_out = td_out.select(*td_module.out_keys) # keep only out keys
>>> grads = torch.autograd.grad(td_out, td_in,
... torch.ones_like(td_out))
Utils¶
|
将一个张量向右扩展以匹配另一个张量的形状。 |
|
将一个张量向右扩展以匹配所需的形状。 |
|
测试 input 中 |
|
移除沿指定维度在 key 中重复的索引。 |
|
获取捕获非张量堆栈的当前设置。 |
|
密集堆叠一个 |
|
|
|
检查数据对象或类型是否是 tensordict 库中的张量容器。 |
|
如果将为选定的方法使用惰性表示,则返回 True。 |
|
从关键字参数或输入字典返回一个创建的 TensorDict。 |
|
合并 tensordicts。 |
|
在批次维度上用常量值填充 tensordict 中的所有张量,并返回一个新的 tensordict。 |
|
填充 tensordicts 列表,以便它们可以被堆叠成连续格式。 |
将 TensorDict repr 解析为 TensorDict。 |
|
一个上下文管理器或装饰器,用于控制是否应将相同的非张量数据堆叠为单个 NonTensorData 对象或 NonTensorStack。 |
|
|
设置一些方法的惰性转换行为。 |
|
上下文管理器和装饰器,用于控制 TensorDict 中列表处理的行为。 |
|
检索 TensorDict 中列表到堆栈转换的当前设置。 |