TensorDictBase¶

class tensordict.TensorDictBase¶

TensorDictBase 是 TensorDicts 的抽象父类，TensorDicts 是 torch.Tensor 数据容器。

abs() → T¶: 计算 TensorDict 中每个元素的绝对值。

abs_() → T¶: 就地计算 TensorDict 中每个元素的绝对值。

acos() → T¶: 计算 TensorDict 中每个元素的 acos() 值。

acos_() → T¶: 就地计算 TensorDict 中每个元素的 acos() 值。

add(other: TensorDictBase | torch.Tensor, *, alpha: float | None = None, default: str | CompatibleType | None = None) → TensorDictBase¶

将 other（乘以 alpha）添加到 self 中。

\[\text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i\]

参数:

other (TensorDictBase 或 torch.Tensor) – 要添加到 self 中的张量或 TensorDict。

关键字参数:

alpha (数字, 可选) – other 的乘数。
default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

add_(other: TensorDictBase | float, *, alpha: float | None = None)¶: add() 的就地版本。

注意

就地 add 不支持 default 关键字参数。

addcdiv(other1: TensorDictBase | torch.Tensor, other2: TensorDictBase | torch.Tensor, value: float | None = 1)¶

执行 other1 除以 other2 的逐元素除法，将结果乘以标量 value，并将结果添加到 self 中。

\[\text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i}\]

self、other1 和 other2 中元素的形状必须可广播。

对于 FloatTensor 或 DoubleTensor 类型的输入，value 必须是实数，否则是整数。

参数:

other1 (TensorDict 或张量) – 分子 tensordict（或张量）
tensor2 (TensorDict 或 Tensor) – 分母张量字典（或张量）

关键字参数:

value (Number, 可选) – \(\text{tensor1} / \text{tensor2}\) 的乘数

addcdiv_(other1, other2, *, value: float | None = 1)¶: addcdiv() 的就地版本。

addcmul(other1, other2, *, value: float | None = 1)¶

对 other1 与 other2 进行逐元素相乘，并将结果乘以标量 value，然后将其加到 self 中。

\[\text{out}_i = \text{input}_i + \text{value} \times \text{other1}_i \times \text{other2}_i\]

self、other1 和 other2 的形状必须是可广播的。

对于 FloatTensor 或 DoubleTensor 类型的输入，value 必须是实数，否则是整数。

参数:

other1 (TensorDict 或 Tensor) – 要相乘的张量字典或张量
other2 (TensorDict 或 Tensor) – 要相乘的张量字典或张量

关键字参数:

value (Number, 可选) – \(other1 .* other2\) 的乘数

addcmul_(other1, other2, *, value: float | None = 1)¶: addcmul() 的就地版本。

抽象 all(dim: int = None) → bool | TensorDictBase¶

检查张量字典中的所有值是否为 True/非空。

参数:: dim (int, 可选) – 如果为 None，则返回一个布尔值，表示所有张量是否都返回 tensor.all() == True。如果为整数，则仅当该维度与张量字典的形状兼容时，才会在指定的维度上调用 all。

抽象 any(dim: int = None) → bool | TensorDictBase¶

检查张量字典中是否有任何值为 True/非空。

参数:: dim (int, 可选) – 如果为 None，则返回一个布尔值，表示所有张量是否都返回 tensor.any() == True。如果为整数，则仅当该维度与张量字典的形状兼容时，才会在指定的维度上调用 all。

apply(fn: Callable, *others: T, batch_size: Sequence[int] | None = None, device: torch.device | None = _NoDefault.ZERO, names: Sequence[str] | None = _NoDefault.ZERO, inplace: bool = False, default: Any = _NoDefault.ZERO, filter_empty: bool | None = None, propagate_lock: bool = False, call_on_nested: bool = False, out: TensorDictBase | None = None, **constructor_kwargs) → T | None¶

将一个可调用对象应用到张量字典中存储的所有值上，并将它们设置到一个新的张量字典中。

可调用对象的签名必须为 Callable[Tuple[Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。

参数:

fn (Callable) – 要应用到张量字典中的张量的函数。
*others (TensorDictBase 实例, 可选) – 如果提供，这些 tensordict 实例应该具有与 self 相匹配的结构。 fn 参数应该接收与 tensordict 数量相同的未命名输入，包括 self。如果其他 tensordict 缺少条目，可以通过 default 关键字参数传递默认值。

关键字参数:

batch_size (int 序列, 可选) – 如果提供，则生成的 TensorDict 将具有所需的 batch_size。 batch_size 参数应该与转换后的 batch_size 相匹配。这是一个关键字参数。
device (torch.device, 可选) – 生成结果的设备，如果有。
names (字符串列表, 可选) – 新的维度名称，如果修改了 batch_size。
inplace (bool, 可选) – 如果为 True，则更改将就地进行。默认值为 False。这是一个关键字参数。
default (任何类型, 可选) – 其他 tensordict 中缺少条目的默认值。如果未提供，则缺少的条目将引发 KeyError。
filter_empty (bool, 可选) – 如果 True，则空 tensordict 将被过滤掉。这也有助于降低计算成本，因为空数据结构不会被创建和销毁。非张量数据被视为叶节点，因此即使函数未对其进行操作，也会保留在 tensordict 中。为确保向后兼容性，默认值为 False。
propagate_lock (bool, 可选) – 如果 True，则已锁定 tensordict 将生成另一个已锁定 tensordict。默认值为 False。

call_on_nested (bool, 可选) –

如果 True，则该函数将在第一级张量和容器（TensorDict 或张量类）上调用。在这种情况下，func 负责将调用传播到嵌套级别。这允许在将调用传播到嵌套 tensordict 时实现细粒度行为。如果 False，则该函数仅在叶节点上调用，并且 apply 将负责将该函数分派到所有叶节点。

>>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]})
>>> def mean_tensor_only(val):
...     if is_tensor_collection(val):
...         raise RuntimeError("Unexpected!")
...     return val.mean()
>>> td_mean = td.apply(mean_tensor_only)
>>> def mean_any(val):
...     if is_tensor_collection(val):
...         # Recurse
...         return val.apply(mean_any, call_on_nested=True)
...     return val.mean()
>>> td_mean = td.apply(mean_any, call_on_nested=True)

out (TensorDictBase, 可选) –
一个 tensordict，用于写入结果。这可以用来避免创建新的 tensordict
```
>>> td = TensorDict({"a": 0})
>>> td.apply(lambda x: x+1, out=td)
>>> assert (td==1).all()
```
警告

如果 tensordict 上执行的操作需要访问多个键才能进行单个计算，则提供一个等于 self 的 out 参数会导致该操作静默地提供错误的结果。例如
```
>>> td = TensorDict({"a": 1, "b": 1})
>>> td.apply(lambda x: x+td["a"])["b"] # Right!
tensor(2)
>>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong!
tensor(3)
```
**constructor_kwargs** – 传递给 TensorDict 构造函数的额外关键字参数。

返回值：:

一个新的 tensordict，包含已转换的张量。

示例

>>> td = TensorDict({ ... "a": -torch.ones(3), ... "b": {"c": torch.ones(3)}}, ... batch_size=[3]) >>> td_1 = td.apply(lambda x: x+1) >>> assert (td_1["a"] == 0).all() >>> assert (td_1["b", "c"] == 2).all() >>> td_2 = td.apply(lambda x, y: x+y, td) >>> assert (td_2["a"] == -2).all() >>> assert (td_2["b", "c"] == 2).all()

注意

如果函数返回 None，则忽略该条目。这可以用来过滤 tensordict 中的数据

>>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, []) >>> def filter(tensor): ... if tensor == 1: ... return tensor >>> td.apply(filter) TensorDict( fields={ 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)

注意

apply 方法将返回一个 TensorDict 实例，无论输入类型是什么。要保留相同的类型，可以执行以下操作

>>> out = td.clone(False).update(td.apply(...))

apply_(fn: Callable, *others, **kwargs) → T¶

将可调用对象应用于 tensordict 中存储的所有值，并就地重新写入它们。

参数:

fn (Callable) – 要应用到张量字典中的张量的函数。

*others (TensorDictBase 序列, 可选) – 将要使用的其他 tensordict。

关键字参数：请参见 apply()。

返回值：:

self 或应用了该函数的 self 的副本

asin() → T¶

计算 TensorDict 中每个元素的 asin() 值。

asin_() → T¶

就地计算 TensorDict 中每个元素的 asin() 值。

atan() → T¶

计算 TensorDict 中每个元素的 atan() 值。

atan_() → T¶

就地计算 TensorDict 中每个元素的 atan() 值。

auto_batch_size_(batch_dims: int | None = None) → T¶

设置 tensordict 的最大 batch-size，最多到可选的 batch_dims。

参数:

batch_dims (int, 可选) – 如果提供，则 batch-size 最多为 batch_dims 长。

返回值：:

self

示例

>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict({"a": torch.randn(3, 4, 5), "b": {"c": torch.randn(3, 4, 6)}}, batch_size=[]) >>> td.auto_batch_size_() >>> print(td.batch_size) torch.Size([3, 4]) >>> td.auto_batch_size_(batch_dims=1) >>> print(td.batch_size) torch.Size([3])

property batch_dims: int¶

tensordict batch 大小的长度。

返回值：:

描述 tensordict 维度数量的整数。

abstract property batch_size: Size¶

TensorDict 的形状（或 batch_size）。

tensordict 的形状对应于它包含的张量的共同前 N 个维度，其中 N 是任意数量。 batch-size 与“特征大小”形成对比，后者代表张量的语义相关形状。例如，一批视频可能具有形状 [B, T, C, W, H]，其中 [B, T] 是 batch-size（批次和时间维度），而 [C, W, H] 是特征维度（通道和空间维度）。

用户在初始化时控制 TensorDict 的形状（即，它不是从张量形状推断出来的）。

如果新大小与 TensorDict 内容兼容，则可以动态编辑 batch_size。例如，将 batch 大小设置为空值始终是允许的。

返回值：:

一个描述 TensorDict batch 大小的 Size 对象。

示例

>>> data = TensorDict({ ... "key 0": torch.randn(3, 4), ... "key 1": torch.randn(3, 5), ... "nested": TensorDict({"key 0": torch.randn(3, 4)}, batch_size=[3, 4])}, ... batch_size=[3]) >>> data.batch_size = () # resets the batch-size to an empty value

bfloat16()¶

将所有张量转换为 torch.bfloat16。

bool()¶

将所有张量转换为 torch.bool。

classmethod cat(input, dim=0, *, out=None)¶

将 tensordict 沿给定维度连接成单个 tensordict。

此调用等效于调用 torch.cat()，但与 torch.compile 兼容。

ceil() → T¶

计算 TensorDict 中每个元素的 ceil() 值。

ceil_() → T¶

对 TensorDict 中每个元素就地计算 ceil() 值。

chunk(chunks: int, dim: int = 0) → tuple[TensorDictBase, ...]¶

如果可能，将一个 tensordict 分割成指定数量的块。

每个块都是输入 tensordict 的视图。

参数:

chunks (int) – 要返回的块的数量

dim (int, optional) – 沿其分割 tensordict 的维度。默认值为 0。

示例

>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 4, 2), ... }, batch_size=[3, 4]) >>> td0, td1 = td.chunk(dim=-1, chunks=2) >>> td0['x'] tensor([[[ 0, 1], [ 2, 3]], [[ 8, 9], [10, 11]], [[16, 17], [18, 19]]])

clamp_max(other: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T¶

如果 self 的元素大于该值，则将其钳制到 other。

参数:

other (TensorDict or Tensor) – 另一个输入 tensordict 或张量。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

clamp_max_(other: TensorDictBase | torch.Tensor) → T¶

clamp_max() 的就地版本。

注意

就地 clamp_max 不支持 default 关键字参数。

clamp_min(other: TensorDictBase | torch.Tensor, default: str | CompatibleType | None = None) → T¶

如果 self 的元素小于该值，则将其钳制到 other。

参数:

other (TensorDict or Tensor) – 另一个输入 tensordict 或张量。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

clamp_min_(other: TensorDictBase | torch.Tensor) → T¶

clamp_min() 的就地版本。

注意

就地 clamp_min 不支持 default 关键字参数。

clear() → T¶

清除 tensordict 的内容。

clear_device_() → T¶

清除 tensordict 的设备。

返回：自身

clone(recurse: bool = True, **kwargs) → T¶

将 TensorDictBase 子类实例克隆到同一类型的新 TensorDictBase 子类上。

要从任何其他 TensorDictBase 子类型创建 TensorDict 实例，请改用 to_tensordict() 方法。

参数:

recurse (bool, optional) – 如果为 True，则 TensorDict 中包含的每个张量也将被复制。否则，只有 TensorDict 树结构会被复制。默认为 True。

注意

与许多其他操作（逐点算术、形状操作、……）不同，clone 不会继承原始锁属性。这种设计选择是为了确保克隆可以被创建以进行修改，这是最常见的用法。

complex128()¶

将所有张量转换为 torch.complex128。

complex32()¶

将所有张量转换为 torch.complex32。

complex64()¶

将所有张量转换为 torch.complex64。

consolidate(filename: Path | str | None = None, *, num_threads=0, device: torch.device | None = None, non_blocking: bool = False, inplace: bool = False, return_early: bool = False, use_buffer: bool = False, share_memory: bool = False, pin_memory: bool = False, metadata: bool = False) → None¶

将 tensordict 内容合并到单个存储中，以实现快速序列化。

参数:

filename (Path, optional) – 用于存储 tensordict 的内存映射张量的可选文件路径。

关键字参数:

num_threads (integer, optional) – 用于填充存储的线程数。

device (torch.device, optional) – 存储必须实例化的可选设备。

non_blocking (bool, optional) – 传递给 copy_() 的 non_blocking 参数。

inplace (bool, optional) – 如果为 True，则结果 tensordict 与具有更新值的 self 相同。默认为 False。

return_early (bool, optional) – 如果为 True 且 num_threads>0，则该方法将返回 tensordict 的 future。可以使用 future.result() 查询结果 tensordict。

use_buffer (bool, optional) – 如果为 True 且传递了 filename，则将在共享内存中创建一个中间本地缓冲区，并将数据作为最后一步复制到存储位置。这可能比直接写入远程物理内存（例如，NFS）更快。默认为 False。

share_memory (bool, optional) – 如果为 True，则存储将放置在共享内存中。默认为 False。

pin_memory (bool, optional) – 是否应将合并后的数据放置在固定内存中。默认为 False。

metadata (bool, optional) – 如果为 True，则元数据将与公共存储一起存储。如果提供了 filename，则此操作无效。存储元数据在需要控制序列化方式时非常有用，因为如果元数据存在或不存在，TensorDict 将以不同的方式处理合并的 TD 的序列化/反序列化。

注意

如果 tensordict 已经合并，则所有参数都将被忽略，并将返回 self。调用 contiguous() 以重新合并。

示例

>>> import pickle >>> import tempfile >>> import torch >>> import tqdm >>> from torch.utils.benchmark import Timer >>> from tensordict import TensorDict >>> data = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}}) >>> data_consolidated = data.consolidate() >>> # check that the data has a single data_ptr() >>> assert torch.tensor([ ... v.untyped_storage().data_ptr() for v in data_c.values(True, True) ... ]).unique().numel() == 1 >>> # Serializing the tensordict will be faster with data_consolidated >>> with open("data.pickle", "wb") as f: ... print("regular", Timer("pickle.dump(data, f)", globals=globals()).adaptive_autorange()) >>> with open("data_c.pickle", "wb") as f: ... print("consolidated", Timer("pickle.dump(data_consolidated, f)", globals=globals()).adaptive_autorange())

abstract contiguous() → T¶

返回一个具有连续值的新 tensordict（如果值已经是连续的，则返回 self）。

copy()¶

返回 tensordict 的浅层副本（即复制结构，但不复制数据）。

等同于 TensorDictBase.clone(recurse=False)

copy_(tensordict: T, non_blocking: bool = False) → T¶

参见 TensorDictBase.update_.

非阻塞参数将被忽略，仅出于与 torch.Tensor.copy_() 的兼容性而存在。

copy_at_(tensordict: T, idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], non_blocking: bool = False) → T¶

参见 TensorDictBase.update_at_.

cos() → T¶

计算 TensorDict 中每个元素的 cos() 值。

cos_() → T¶

就地计算 TensorDict 中每个元素的 cos() 值。

cosh() → T¶

计算 TensorDict 中每个元素的 cosh() 值。

cosh_() → T¶

就地计算 TensorDict 中每个元素的 cosh() 值。

cpu(**kwargs) → T¶

将 TensorDict 转换为 CPU。

此函数还支持 to() 的所有关键字参数。

create_nested(key)¶

创建一个嵌套的 TensorDict，其形状、设备和维度名称与当前 TensorDict 相同。

如果该值已存在，则此操作将覆盖它。此操作在锁定的 TensorDict 中被阻止。

示例

>>> data = TensorDict({}, [3, 4, 5]) >>> data.create_nested("root") >>> data.create_nested(("some", "nested", "value")) >>> print(data) TensorDict( fields={ root: TensorDict( fields={ }, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False), some: TensorDict( fields={ nested: TensorDict( fields={ value: TensorDict( fields={ }, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)

cuda(device: Optional[int] = None, **kwargs) → T¶

将 TensorDict 转换为 CUDA 设备（如果尚未在该设备上）。

参数:

device (int, optional) – 如果提供，则为应将张量转换为其上的 CUDA 设备。

此函数还支持 to() 的所有关键字参数。

property data¶

返回一个包含叶张量的 .data 属性的 TensorDict。

abstract del_(key: NestedKey) → T¶

删除 TensorDict 的一个键。

参数:

key (NestedKey) – 要删除的键

返回值：:

self

property depth: int¶

返回 TensorDict 的深度 - 最大级别数。

最小深度为 0（没有嵌套 TensorDict）。

detach() → T¶

分离 TensorDict 中的张量。

返回值：:

一个没有张量需要梯度的新的 TensorDict。

abstract detach_() → T¶

就地分离 TensorDict 中的张量。

返回值：:

self.

abstract property device: torch.device | None¶

TensorDict 的设备。

如果 TensorDict 有一个指定的设备，则它的所有张量（包括嵌套的张量）都必须位于同一个设备上。如果 TensorDict 设备为 None，则不同的值可以位于不同的设备上。

返回值：:

torch.device 对象，指示张量放置的设备，如果 TensorDict 没有设备，则为 None。

示例

>>> td = TensorDict({ ... "cpu": torch.randn(3, device='cpu'), ... "cuda": torch.randn(3, device='cuda'), ... }, batch_size=[], device=None) >>> td['cpu'].device device(type='cpu') >>> td['cuda'].device device(type='cuda') >>> td = TensorDict({ ... "x": torch.randn(3, device='cpu'), ... "y": torch.randn(3, device='cuda'), ... }, batch_size=[], device='cuda') >>> td['x'].device device(type='cuda') >>> td['y'].device device(type='cuda') >>> td = TensorDict({ ... "x": torch.randn(3, device='cpu'), ... "y": TensorDict({'z': torch.randn(3, device='cpu')}, batch_size=[], device=None), ... }, batch_size=[], device='cuda') >>> td['x'].device device(type='cuda') >>> td['y'].device # nested tensordicts are also mapped onto the appropriate device. device(type='cuda') >>> td['y', 'x'].device device(type='cuda')

dim() → int¶

参见 batch_dims().

div(other: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T¶

将输入 self 的每个元素除以 other 的对应元素。

\[\text{out}_i = \frac{\text{input}_i}{\text{other}_i}\]

支持广播、类型提升以及整数、浮点数、TensorDict 或张量输入。始终将整数类型提升为默认标量类型。

参数:

other (TensorDict, Tensor 或 Number) – 除数。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

div_(other: TensorDictBase | torch.Tensor) → T¶

div() 的原地版本。

注意

原地 div 不支持 default 关键字参数。

double()¶

将所有张量转换为 torch.bool。

property dtype¶

如果张量字典中的值具有唯一的 dtype，则返回该 dtype。

dumps(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

将张量字典保存到磁盘。

此函数是 memmap() 的代理。

empty(recurse=False, *, batch_size=None, device=_NoDefault.ZERO, names=None) → T¶

返回一个新的、空的张量字典，具有相同的设备和批次大小。

参数:

recurse (bool, 可选) – 如果为 True，则将复制 TensorDict 的整个结构，不含内容。否则，只会复制根。默认为 False。

关键字参数:

batch_size (torch.Size, 可选) – 张量字典的新批次大小。

device (torch.device, 可选) – 新设备。

names (str 列表, 可选) – 维度名称。

abstract entry_class(key: NestedKey) → type¶

返回条目的类，可能避免调用 isinstance(td.get(key), type)。

此方法应优先于 tensordict.get(key).shape，只要 get() 的执行成本很高。

erf() → T¶

计算张量字典中每个元素的 erf() 值。

erf_() → T¶

原地计算张量字典中每个元素的 erf() 值。

erfc() → T¶

计算张量字典中每个元素的 erfc() 值。

erfc_() → T¶

原地计算张量字典中每个元素的 erfc() 值。

exclude(*keys: NestedKey, inplace: bool = False) → T¶

排除张量字典的键，并返回一个不包含这些条目的新张量字典。

不会复制值：对原始张量字典或新张量字典中任何一个张量的原地修改都会导致两个张量字典都发生变化。

参数:

*keys (str) – 要排除的键。

inplace (bool) – 如果为 True，则原地修剪张量字典。默认值为 False。

返回值：:

一个不包含被排除条目的新张量字典（如果 inplace=True，则为同一个张量字典）。

示例

>>> from tensordict import TensorDict >>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, []) >>> td.exclude("a", ("b", "c")) TensorDict( fields={ b: TensorDict( fields={ d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.exclude("a", "b") TensorDict( fields={ }, batch_size=torch.Size([]), device=None, is_shared=False)

exp() → T¶

计算张量字典中每个元素的 exp() 值。

exp_() → T¶

原地计算张量字典中每个元素的 exp() 值。

abstract expand(*shape: int) → T¶

abstract expand(shape: Size) → T

根据 expand() 函数扩展张量字典中的每个张量，忽略特征维度。

支持使用可迭代对象指定形状。

示例

>>> td = TensorDict({ ... 'a': torch.zeros(3, 4, 5), ... 'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4]) >>> td_expand = td.expand(10, 3, 4) >>> assert td_expand.shape == torch.Size([10, 3, 4]) >>> assert td_expand.get("a").shape == torch.Size([10, 3, 4, 5])

expand_as(other: TensorDictBase | torch.Tensor) → TensorDictBase¶

将张量字典的形状广播到 other 的形状，并相应地扩展它。

如果输入是张量集合（tensordict 或 tensorclass），则叶子将以一对一的方式扩展。

示例

>>> from tensordict import TensorDict >>> import torch >>> td0 = TensorDict({ ... "a": torch.ones(3, 1, 4), ... "b": {"c": torch.ones(3, 2, 1, 4)}}, ... batch_size=[3], ... ) >>> td1 = TensorDict({ ... "a": torch.zeros(2, 3, 5, 4), ... "b": {"c": torch.zeros(2, 3, 2, 6, 4)}}, ... batch_size=[2, 3], ... ) >>> expanded = td0.expand_as(td1) >>> assert (expanded==1).all() >>> print(expanded) TensorDict( fields={ a: Tensor(shape=torch.Size([2, 3, 5, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([2, 3, 2, 6, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([2, 3]), device=None, is_shared=False)}, batch_size=torch.Size([2, 3]), device=None, is_shared=False)

expm1() → T¶

计算 TensorDict 中每个元素的 expm1() 值。

expm1_() → T¶

就地计算 TensorDict 中每个元素的 expm1() 值。

fill_(key: NestedKey, value: float | bool) → T¶

用给定的标量值填充由键指向的张量。

参数:

key (str or nested key) – 要填充的条目。

value (Number or bool) – 用于填充的值。

返回值：:

self

filter_empty_()¶

就地过滤掉所有空的 tensordicts。

filter_non_tensor_data() → T¶

过滤掉所有非张量数据。

flatten(start_dim=0, end_dim=- 1)¶

展平 tensordict 中的所有张量。

参数:

start_dim (int) – 要展平的第一个维度

end_dim (int) – 要展平的最后一个维度

示例

>>> td = TensorDict({ ... "a": torch.arange(60).view(3, 4, 5), ... "b": torch.arange(12).view(3, 4)}, batch_size=[3, 4]) >>> td_flat = td.flatten(0, 1) >>> td_flat.batch_size torch.Size([12]) >>> td_flat["a"] 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]]) >>> td_flat["b"] tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])

flatten_keys(separator: str = '.', inplace: bool = False, is_leaf: Callable[[Type], bool] | None = None) → T¶

递归地将嵌套的 tensordict 转换为扁平的 tensordict。

TensorDict 类型将丢失，结果将是一个简单的 TensorDict 实例。

参数:

separator (str, optional) – 嵌套项之间的分隔符。

inplace (bool, optional) – 如果为 True，则生成的 tensordict 将与调用发生所在的 tensordict 具有相同的标识。默认为 False。

is_leaf (callable, optional) – 一个可调用对象，用于对类类型进行操作，返回一个布尔值，指示该类是否必须被视为叶子。

示例

>>> data = TensorDict({"a": 1, ("b", "c"): 2, ("e", "f", "g"): 3}, batch_size=[]) >>> data.flatten_keys(separator=" - ") TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b - c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), e - f - g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)

此方法和 unflatten_keys() 在处理状态字典时特别有用，因为它们使您可以无缝地将扁平字典转换为模拟模型结构的数据结构。

示例

>>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4)) >>> ddp_model = torch.ao.quantization.QuantWrapper(model) >>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".") >>> print(state_dict) TensorDict( fields={ module: TensorDict( fields={ 0: TensorDict( fields={ bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> model_state_dict = state_dict.get("module") >>> print(model_state_dict) TensorDict( fields={ 0: TensorDict( fields={ bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> model.load_state_dict(dict(model_state_dict.flatten_keys(".")))

float()¶

将所有张量转换为 torch.float。

float16()¶

将所有张量转换为 torch.float16。

float32()¶

将所有张量转换为 torch.float32。

float64()¶

将所有张量转换为 torch.float64。

floor() → T¶

计算 TensorDict 中每个元素的 floor() 值。

floor_() → T¶

就地计算 TensorDict 中每个元素的 floor() 值。

frac() → T¶

计算 TensorDict 中每个元素的 frac() 值。

frac_() → T¶

就地计算 TensorDict 中每个元素的 frac() 值。

abstract classmethod from_dict(input_dict, batch_size: torch.Size | None = None, device: torch.device | None = None, batch_dims: int | None = None, names: List[str] | None = None)¶

返回一个从字典或其他 TensorDict 创建的 TensorDict。

如果未指定 batch_size，则返回可能的最大批次大小。

此函数也适用于嵌套字典，或可用于确定嵌套 tensordict 的批次大小。

参数:

input_dict (dictionary, optional) – 用作数据源的字典（嵌套键兼容）。

batch_size (iterable of int, optional) – tensordict 的批次大小。

device (torch.device or compatible type, optional) – TensorDict 的设备。

batch_dims (int, 可选) – batch_dims（即被视为batch_size 的前导维度数量）。与 batch_size 互斥。请注意，这是 tensordict 的__最大__批次维度数量，允许更小的数量。

names (字符串列表, 可选) – tensordict 的维度名称。

示例

>>> input_dict = {"a": torch.randn(3, 4), "b": torch.randn(3)} >>> print(TensorDict.from_dict(input_dict)) TensorDict( fields={ a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False) >>> # nested dict: the nested TensorDict can have a different batch-size >>> # as long as its leading dims match. >>> input_dict = {"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}} >>> print(TensorDict.from_dict(input_dict)) TensorDict( fields={ a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False) >>> # we can also use this to work out the batch sie of a tensordict >>> input_td = TensorDict({"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}, []) >>> print(TensorDict.from_dict(input_td)) TensorDict( fields={ a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False)

抽象 from_dict_instance(input_dict, batch_size=None, device=None, batch_dims=None, names: List[str] | None = None)¶

from_dict() 的实例方法版本。

与 from_dict() 不同，此方法将尝试在现有树内（对于任何现有叶子）保持 tensordict 类型。

示例

>>> from tensordict import TensorDict, tensorclass >>> import torch >>> >>> @tensorclass >>> class MyClass: ... x: torch.Tensor ... y: int >>> >>> td = TensorDict({"a": torch.randn(()), "b": MyClass(x=torch.zeros(()), y=1)}) >>> print(td.from_dict_instance(td.to_dict())) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: MyClass( x=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), y=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> print(td.from_dict(td.to_dict())) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), y: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)

类方法 from_h5(filename, mode='r')¶

从 h5 文件创建 PersistentTensorDict。

此函数将自动确定每个嵌套 tensordict 的批次大小。

参数:

filename (字符串) – h5 文件的路径。

mode (字符串, 可选) – 读取模式。默认为 "r"。

类方法 from_module(module, as_module: bool = False, lock: bool = True, use_state_dict: bool = False)¶

将模块的参数和缓冲区复制到 tensordict 中。

参数:

module (nn.Module) – 获取参数的模块。

as_module (bool, 可选) – 如果 True，将返回一个 TensorDictParams 实例，可用于存储 torch.nn.Module 中的参数。默认为 False。

lock (bool, 可选) – 如果 True，生成的 tensordict 将被锁定。默认为 True。

use_state_dict (bool, 可选) –
如果 True，将使用来自模块的状态字典，并将状态字典解扁平化为具有模型树结构的 TensorDict。默认为 False。 .. 注意

This is particularly useful when state-dict hooks have to be used.

示例

>>> from torch import nn >>> module = nn.TransformerDecoder( ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), ... num_layers=1) >>> params = TensorDict.from_module(module) >>> print(params["layers", "0", "linear1"]) TensorDict( fields={ bias: Parameter(shape=torch.Size([2048]), device=cpu, dtype=torch.float32, is_shared=False), weight: Parameter(shape=torch.Size([2048, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)

类方法 from_modules(*modules, as_module: bool = False, lock: bool = True, use_state_dict: bool = False, lazy_stack: bool = False, expand_identical: bool = False)¶

通过 vmap 获取多个模块的参数，用于集成学习/特征的预期应用程序。

参数:

modules (nn.Module 序列) – 获取参数的模块。如果模块的结构不同，则需要延迟堆叠（请参阅下面的 lazy_stack 参数）。

关键字参数:

as_module (bool, 可选) – 如果 True，将返回一个 TensorDictParams 实例，可用于存储 torch.nn.Module 中的参数。默认为 False。

lock (bool, 可选) – 如果 True，生成的 tensordict 将被锁定。默认为 True。

use_state_dict (bool, 可选) –
如果 True，将使用来自模块的状态字典，并将状态字典解扁平化为具有模型树结构的 TensorDict。默认为 False。 .. 注意

This is particularly useful when state-dict hooks have to be used.

lazy_stack (bool, 可选) –
参数是否应该被密集堆叠或延迟堆叠。默认为 False（密集堆叠）。

注意

lazy_stack 和 as_module 是互斥的特性。

警告

延迟输出和非延迟输出之间存在一个关键区别，即非延迟输出将使用所需的批次大小重新实例化参数，而 lazy_stack 只是将参数表示为延迟堆叠。这意味着，虽然当 lazy_stack=True 时，原始参数可以安全地传递给优化器，但当它设置为 True 时，需要传递新参数。

警告

虽然使用延迟堆叠来保留原始参数引用可能很诱人，但请记住，每次调用 get() 时，延迟堆叠都会执行堆叠。这将需要内存（参数大小的 N 倍，如果构建图形则更多）以及计算时间。这也意味着优化器将包含更多参数，并且像 step() 或 zero_grad() 这样的操作将需要更长的时间才能执行。一般来说，lazy_stack 应该保留在极少数用例中。

expand_identical (bool, 可选) – 如果 True 并且相同参数（相同标识）被堆叠到自身，则将返回该参数的扩展版本。当 lazy_stack=True 时，此参数将被忽略。

示例

>>> from torch import nn >>> from tensordict import TensorDict >>> torch.manual_seed(0) >>> empty_module = nn.Linear(3, 4, device="meta") >>> n_models = 2 >>> modules = [nn.Linear(3, 4) for _ in range(n_models)] >>> params = TensorDict.from_modules(*modules) >>> print(params) TensorDict( fields={ bias: Parameter(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Parameter(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([2]), device=None, is_shared=False) >>> # example of batch execution >>> def exec_module(params, x): ... with params.to_module(empty_module): ... return empty_module(x) >>> x = torch.randn(3) >>> y = torch.vmap(exec_module, (0, None))(params, x) >>> assert y.shape == (n_models, 4) >>> # since lazy_stack = False, backprop leaves the original params untouched >>> y.sum().backward() >>> assert params["weight"].grad.norm() > 0 >>> assert modules[0].weight.grad is None

使用 lazy_stack=True 时，情况略有不同

>>> params = TensorDict.from_modules(*modules, lazy_stack=True) >>> print(params) LazyStackedTensorDict( fields={ bias: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, exclusive_fields={ }, batch_size=torch.Size([2]), device=None, is_shared=False, stack_dim=0) >>> # example of batch execution >>> y = torch.vmap(exec_module, (0, None))(params, x) >>> assert y.shape == (n_models, 4) >>> y.sum().backward() >>> assert modules[0].weight.grad is not None

类方法 from_namedtuple(named_tuple, *, auto_batch_size: bool = False)¶

递归地将命名元组转换为 TensorDict。

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，则自动计算批次大小。默认为 False。

示例

>>> from tensordict import TensorDict >>> import torch >>> data = TensorDict({ ... "a_tensor": torch.zeros((3)), ... "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3]) >>> nt = data.to_namedtuple() >>> print(nt) GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!')) >>> TensorDict.from_namedtuple(nt, auto_batch_size=True) TensorDict( fields={ a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), nested: TensorDict( fields={ a_string: NonTensorData(data=zero!, batch_size=torch.Size([3]), device=None), a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False)

classmethod from_pytree(pytree, *, batch_size: torch.Size | None = None, auto_batch_size: bool = False, batch_dims: int | None = None)¶

将 pytree 转换为 TensorDict 实例。

此方法旨在尽可能保持 pytree 的嵌套结构。

添加了额外的非张量键以跟踪每个级别的标识，从而提供了一个内置的 pytree-to-tensordict 双射转换 API。

目前接受的类包括列表、元组、命名元组和字典。

注意

对于字典，非 NestedKey 键将作为 NonTensorData 实例单独注册。

注意

可转换为张量的类型（例如 int、float 或 np.ndarray）将转换为 torch.Tensor 实例。注意，此转换是满射的：将 tensordict 转换回 pytree 不会恢复原始类型。

示例

>>> # Create a pytree with tensor leaves, and one "weird"-looking dict key >>> class WeirdLookingClass: ... pass ... >>> weird_key = WeirdLookingClass() >>> # Make a pytree with tuple, lists, dict and namedtuple >>> pytree = ( ... [torch.randint(10, (3,)), torch.zeros(2)], ... { ... "tensor": torch.randn( ... 2, ... ), ... "td": TensorDict({"one": 1}), ... weird_key: torch.randint(10, (2,)), ... "list": [1, 2, 3], ... }, ... {"named_tuple": TensorDict({"two": torch.ones(1) * 2}).to_namedtuple()}, ... ) >>> # Build a TensorDict from that pytree >>> td = TensorDict.from_pytree(pytree) >>> # Recover the pytree >>> pytree_recon = td.to_pytree() >>> # Check that the leaves match >>> def check(v1, v2): >>> assert (v1 == v2).all() >>> >>> torch.utils._pytree.tree_map(check, pytree, pytree_recon) >>> assert weird_key in pytree_recon[1]

classmethod fromkeys(keys: List[NestedKey], value: Any = 0)¶

根据键列表和单个值创建 tensordict。

参数:

keys (list of NestedKey) – 指定新字典键的迭代器。

value (兼容类型, optional) – 所有键的值。默认为 0。

gather(dim: int, index: Tensor, out: T | None = None) → T¶

沿由 dim 指定的轴收集值。

参数:

dim (int) – 收集元素的维度

index (torch.Tensor) – 一个长张量，其维度数与 tensordict 相匹配，只有两个之间的一个维度不同（收集维度）。它的元素指的是要沿所需维度收集的索引。

out (TensorDictBase, optional) – 目标 tensordict。它必须与索引具有相同的形状。

示例

>>> td = TensorDict( ... {"a": torch.randn(3, 4, 5), ... "b": TensorDict({"c": torch.zeros(3, 4, 5)}, [3, 4, 5])}, ... [3, 4]) >>> index = torch.randint(4, (3, 2)) >>> td_gather = td.gather(dim=1, index=index) >>> print(td_gather) TensorDict( fields={ a: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 2, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 2]), device=None, is_shared=False)

Gather 保留维度名称。

示例

>>> td.names = ["a", "b"] >>> td_gather = td.gather(dim=1, index=index) >>> td_gather.names ["a", "b"]

gather_and_stack(dst: int, group: 'dist.ProcessGroup' | None = None) → T | None¶

从各个 worker 收集 tensordict，并将它们堆叠到目标 worker 的 self 上。

参数:

dst (int) – 目标 worker 的秩，将在其上调用 gather_and_stack()。

group (torch.distributed.ProcessGroup, optional) – 如果设置，则使用指定的进程组进行通信。否则，将使用默认进程组。默认为 None。

示例

>>> from torch import multiprocessing as mp >>> from tensordict import TensorDict >>> import torch >>> >>> def client(): ... torch.distributed.init_process_group( ... "gloo", ... rank=1, ... world_size=2, ... init_method=f"tcp://127.0.0.1:10003", ... ) ... # Create a single tensordict to be sent to server ... td = TensorDict( ... {("a", "b"): torch.randn(2), ... "c": torch.randn(2)}, [2] ... ) ... td.gather_and_stack(0) ... >>> def server(): ... torch.distributed.init_process_group( ... "gloo", ... rank=0, ... world_size=2, ... init_method=f"tcp://127.0.0.1:10003", ... ) ... # Creates the destination tensordict on server. ... # The first dim must be equal to world_size-1 ... td = TensorDict( ... {("a", "b"): torch.zeros(2), ... "c": torch.zeros(2)}, [2] ... ).expand(1, 2).contiguous() ... td.gather_and_stack(0) ... assert td["a", "b"] != 0 ... print("yuppie") ... >>> if __name__ == "__main__": ... mp.set_start_method("spawn") ... ... main_worker = mp.Process(target=server) ... secondary_worker = mp.Process(target=client) ... ... main_worker.start() ... secondary_worker.start() ... ... main_worker.join() ... secondary_worker.join()

get(key: NestedKey, default: Any = _NoDefault.ZERO) → Tensor¶

获取使用输入键存储的值。

参数:

key (str, tuple of str) – 要查询的键。如果为 str 元组，则等效于 getattr 的链式调用。

default – 如果 tensordict 中没有找到该键，则为默认值。

示例

>>> td = TensorDict({"x": 1}, batch_size=[]) >>> td.get("x") tensor(1) >>> td.get("y", default=None) None

get_at(key: NestedKey, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], default: Tensor = _NoDefault.ZERO) → Tensor¶

从键 key 处的索引 idx 获取 tensordict 的值。

参数:

key (str, tuple of str) – 要检索的键。

索引 (int, 切片, torch.Tensor, 可迭代对象) – 张量的索引。

默认值 (torch.Tensor) – 如果键不存在于张量字典中，则返回的默认值。

返回值：:

索引张量。

示例

>>> td = TensorDict({"x": torch.arange(3)}, batch_size=[]) >>> td.get_at("x", index=1) tensor(1)

get_item_shape(key: NestedKey)¶

返回条目的大小，可能避免重复使用 get().

get_non_tensor(key: NestedKey, default=_NoDefault.ZERO)¶

获取非张量值（如果存在），或者获取 default 值（如果找不到非张量值）。

此方法对张量/TensorDict 值很健壮，这意味着如果收集到的值为普通张量，它也会被返回（尽管此方法有一些开销，不应在超出其自然范围的情况下使用）。

参见 set_non_tensor()，以了解如何在张量字典中设置非张量值。

参数:

key (NestedKey) – NonTensorData 对象的位置。

default (Any, 可选) – 如果找不到键，则返回的值。

返回：tensordict.tensorclass.NonTensorData 的内容，
或对应于 key 的条目（如果不是 tensordict.tensorclass.NonTensorData）（或者找不到条目则返回 default）。

示例

>>> data = TensorDict({}, batch_size=[]) >>> data.set_non_tensor(("nested", "the string"), "a string!") >>> assert data.get_non_tensor(("nested", "the string")) == "a string!" >>> # regular `get` works but returns a NonTensorData object >>> data.get(("nested", "the string")) NonTensorData( data='a string!', batch_size=torch.Size([]), device=None, is_shared=False)

property grad¶

返回一个包含叶张量 .grad 属性的张量字典。

half()¶

将所有张量转换为 torch.half。

int()¶

将所有张量转换为 torch.int。

int16()¶

将所有张量转换为 torch.int16。

int32()¶

将所有张量转换为 torch.int32。

int64()¶

将所有张量转换为 torch.int64。

int8()¶

将所有张量转换为 torch.int8。

irecv(src: int, *, group: 'dist.ProcessGroup' | None = None, return_premature: bool = False, init_tag: int = 0, pseudo_rand: bool = False) → tuple[int, list[torch.Future]] | list[torch.Future] | None¶

异步接收张量字典的内容并用它更新内容。

有关上下文，请查看 isend() 方法中的示例。

参数:

src (int) – 源工作者的等级。

关键字参数:

group (torch.distributed.ProcessGroup, optional) – 如果设置，则使用指定的进程组进行通信。否则，将使用默认进程组。默认为 None。

return_premature (bool) – 如果为 True，则返回一个 futures 列表，在张量字典更新之前等待它。默认为 False，即等待在调用中完成更新。

init_tag (int) – 源工作者使用的 init_tag。

pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许从不同的节点发送多个数据而不重叠。注意，这些伪随机数的生成是昂贵的（1e-5 秒/数），这意味着它可能会减慢算法的运行时间。此值必须与传递给 isend() 的值匹配。默认为 False。

返回值：:

如果 return_premature=True，则返回一个 futures 列表，在张量字典更新之前等待
它。

is_consolidated()¶

检查张量字典是否具有合并的存储。

abstract is_contiguous() → bool¶

返回一个布尔值，指示所有张量是否都是连续的。

is_empty() → bool¶

检查张量字典是否包含任何叶节点。

is_memmap() → bool¶

检查张量字典是否已内存映射。

如果 TensorDict 实例已内存映射，则它将被锁定（条目无法重命名、删除或添加）。如果使用所有内存映射的张量创建 TensorDict，这并不意味着 is_memmap 将返回 True（因为新张量可能已内存映射，也可能未内存映射）。只有在调用 tensordict.memmap_() 时，张量字典才会被视为内存映射。

对于 CUDA 设备上的张量字典，这始终为 True。

is_shared() → bool¶

检查张量字典是否在共享内存中。

如果 TensorDict 实例位于共享内存中，则它会被锁定（条目不能重命名、删除或添加）。如果使用位于共享内存中的所有张量创建一个 TensorDict，并不意味着 is_shared 会返回 True（因为新的张量可能在共享内存中，也可能不在）。只有在调用 tensordict.share_memory_() 或将 tensordict 放置在默认情况下其内容共享的设备（例如，"cuda"）上时，tensordict 才会被视为位于共享内存中。

对于 CUDA 设备上的张量字典，这始终为 True。

isend(dst: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → int¶

异步发送 tensordict 的内容。

参数:

dst (int) – 目标工作者的等级，内容应发送到该工作者。

关键字参数:

group (torch.distributed.ProcessGroup, optional) – 如果设置，则使用指定的进程组进行通信。否则，将使用默认进程组。默认为 None。

init_tag (int) – 用于标记张量的初始标签。请注意，此标签将根据 TensorDict 中包含的张量数量进行递增。

pseudo_rand (bool) – 如果为 True，标签序列将是伪随机的，允许从不同的节点发送多个数据而不会重叠。请注意，这些伪随机数的生成很昂贵（1e-5 秒/数），这意味着它可能会降低算法的运行时间。默认为 False。

示例

>>> import torch >>> from tensordict import TensorDict >>> from torch import multiprocessing as mp >>> def client(): ... torch.distributed.init_process_group( ... "gloo", ... rank=1, ... world_size=2, ... init_method=f"tcp://127.0.0.1:10003", ... ) ... ... td = TensorDict( ... { ... ("a", "b"): torch.randn(2), ... "c": torch.randn(2, 3), ... "_": torch.ones(2, 1, 5), ... }, ... [2], ... ) ... td.isend(0) ... >>> >>> def server(queue, return_premature=True): ... torch.distributed.init_process_group( ... "gloo", ... rank=0, ... world_size=2, ... init_method=f"tcp://127.0.0.1:10003", ... ) ... td = TensorDict( ... { ... ("a", "b"): torch.zeros(2), ... "c": torch.zeros(2, 3), ... "_": torch.zeros(2, 1, 5), ... }, ... [2], ... ) ... out = td.irecv(1, return_premature=return_premature) ... if return_premature: ... for fut in out: ... fut.wait() ... assert (td != 0).all() ... queue.put("yuppie") ... >>> >>> if __name__ == "__main__": ... queue = mp.Queue(1) ... main_worker = mp.Process( ... target=server, ... args=(queue, ) ... ) ... secondary_worker = mp.Process(target=client) ... ... main_worker.start() ... secondary_worker.start() ... out = queue.get(timeout=10) ... assert out == "yuppie" ... main_worker.join() ... secondary_worker.join()

isfinite() → T¶

返回一个新的 tensordict，其中包含布尔元素，表示每个元素是否有限。

当实数不为 NaN、负无穷大或无穷大时，它们是有限的。当复数的实部和虚部都是有限的时，它们是有限的。

isnan() → T¶

返回一个新的 tensordict，其中包含布尔元素，表示输入的每个元素是否为 NaN。

当复数的实部或虚部为 NaN 时，它们被认为是 NaN。

isneginf() → T¶

测试输入的每个元素是否为负无穷大。

isposinf() → T¶

测试输入的每个元素是否为负无穷大。

isreal() → T¶

返回一个新的 tensordict，其中包含布尔元素，表示输入的每个元素是否为实数。

items(include_nested: bool = False, leaves_only: bool = False, is_leaf=None) → Iterator[tuple[str, CompatibleType]]¶

返回 tensordict 的键值对生成器。

参数:

include_nested (bool, optional) – 如果为 True，将返回嵌套值。默认为 False。

leaves_only (bool, optional) – 如果为 False，将只返回叶子。默认为 False。

is_leaf – 一个可选的 callable，用于指示某个类是否应被视为叶子。

abstract keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None)¶

返回 tensordict 键的生成器。

参数:

include_nested (bool, optional) – 如果为 True，将返回嵌套值。默认为 False。

leaves_only (bool, optional) – 如果为 False，将只返回叶子。默认为 False。

is_leaf – 一个可选的 callable，用于指示某个类是否应被视为叶子。

示例

>>> from tensordict import TensorDict >>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[]) >>> data.keys() ['0', '1'] >>> list(data.keys(leaves_only=True)) ['0'] >>> list(data.keys(include_nested=True, leaves_only=True)) ['0', '1', ('1', '2')]

classmethod lazy_stack(input, dim=0, *, out=None, **kwargs)¶

创建一个 tensordict 的延迟堆栈。

有关详细信息，请参见 lazy_stack()。

lerp(end: TensorDictBase | torch.Tensor, weight: TensorDictBase | torch.Tensor | float)¶

根据标量或张量 weight，对两个张量 start（由 self 给出）和 end 进行线性插值。

\[\text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i)\]

start 和 end 的形状必须是可广播的。如果 weight 是张量，那么 weight、start 和 end 的形状必须是可广播的。

参数:

end (TensorDict) – 包含结束点的 tensordict。

weight (TensorDict, tensor or float) – 插值公式的权重。

lerp_(end: TensorDictBase | float, weight: TensorDictBase | float)¶

lerp() 的就地版本。

lgamma() → T¶

计算 TensorDict 中每个元素的 lgamma() 值。

lgamma_() → T¶

就地计算 TensorDict 中每个元素的 lgamma() 值。

classmethod load(prefix: str | Path, *args, **kwargs) → T¶

从磁盘加载 tensordict。

此类方法是 load_memmap() 的代理。

load_(prefix: str | Path, *args, **kwargs)¶

将磁盘上的 tensordict 加载到当前 tensordict 中。

此类方法是 load_memmap_() 的代理。

classmethod load_memmap(prefix: str | Path, device: torch.device | None = None, non_blocking: bool = False, *, out: TensorDictBase | None = None) → T¶

从磁盘加载内存映射的 tensordict。

参数:

prefix (str 或 Path to folder) – 保存 tensordict 的文件夹的路径。

device (torch.device 或等效, 可选) – 如果提供，数据将异步转换为该设备。支持 “meta” 设备，在这种情况下，数据不会加载，而是会创建一组空的“meta”张量。这有助于了解模型的总大小和结构，而无需实际打开任何文件。

non_blocking (bool, 可选) – 如果为 True，则在设备上加载张量后不会调用同步。默认为 False。

out (TensorDictBase, 可选) – 可选的 tensordict，数据应写入其中。

示例

>>> from tensordict import TensorDict >>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0) >>> td.memmap("./saved_td") >>> td_load = TensorDict.load_memmap("./saved_td") >>> assert (td == td_load).all()

此方法还允许加载嵌套的 tensordict。

>>> nested = TensorDict.load_memmap("./saved_td/nested") >>> assert nested["e"] == 0

tensordict 也可以在“meta”设备上加载，或者作为伪张量加载

>>> import tempfile >>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}}) >>> with tempfile.TemporaryDirectory() as path: ... td.save(path) ... td_load = TensorDict.load_memmap(path, device="meta") ... print("meta:", td_load) ... from torch._subclasses import FakeTensorMode ... with FakeTensorMode(): ... td_load = TensorDict.load_memmap(path) ... print("fake:", td_load) meta: TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=meta, is_shared=False)}, batch_size=torch.Size([]), device=meta, is_shared=False) fake: TensorDict( fields={ a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False)

load_memmap_(prefix: str | Path)¶

在调用 load_memmap_ 的 tensordict 中加载内存映射的 tensordict 的内容。

有关更多信息，请参阅 load_memmap()。

load_state_dict(state_dict: OrderedDict[str, Any], strict=True, assign=False, from_flatten=False) → T¶

将格式为 state_dict() 的状态字典加载到 tensordict 中。

参数:

state_dict (OrderedDict) – 要复制的状态字典。

strict (bool, 可选) – 是否严格要求 state_dict 中的键与此 tensordict 的 torch.nn.Module.state_dict() 函数返回的键匹配。默认：True

assign (bool, 可选) – 是否将状态字典中的项目分配到 tensordict 中的相应键，而不是就地将它们复制到 tensordict 的当前张量中。当为 False 时，当前模块中张量的属性将保留，而当为 True 时，状态字典中张量的属性将保留。默认：False

from_flatten (bool, 可选) – 如果为 True，则假定输入状态字典已扁平化。默认为 False。

示例

>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, []) >>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, []) >>> sd = data.state_dict() >>> data_zeroed.load_state_dict(sd) >>> print(data_zeroed["3", "3"]) tensor(3) >>> # with flattening >>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, []) >>> data_zeroed.load_state_dict(data.state_dict(flatten=True), from_flatten=True) >>> print(data_zeroed["3", "3"]) tensor(3)

lock_() → T¶

将 tensordict 锁定以进行非就地操作。

例如 set()、__setitem__()、update()、rename_key_() 或其他添加或删除条目的操作将被阻止。

此方法可用作装饰器。

示例

>>> from tensordict import TensorDict >>> td = TensorDict({"a": 1, "b": 2, "c": 3}, batch_size=[]) >>> with td.lock_(): ... assert td.is_locked ... try: ... td.set("d", 0) # error! ... except RuntimeError: ... print("td is locked!") ... try: ... del td["d"] ... except RuntimeError: ... print("td is locked!") ... try: ... td.rename_key_("a", "d") ... except RuntimeError: ... print("td is locked!") ... td.set("a", 0, inplace=True) # No storage is added, moved or removed ... td.set_("a", 0) # No storage is added, moved or removed ... td.update({"a": 0}, inplace=True) # No storage is added, moved or removed ... td.update_({"a": 0}) # No storage is added, moved or removed >>> assert not td.is_locked

log() → T¶

计算 TensorDict 中每个元素的 log() 值。

log10() → T¶

计算 TensorDict 中每个元素的 log10() 值。

log10_() → T¶

就地计算 TensorDict 中每个元素的 log10() 值。

log1p() → T¶

计算 TensorDict 中每个元素的 log1p() 值。

log1p_() → T¶

就地计算 TensorDict 中每个元素的 log1p() 值。

log2() → T¶

计算 TensorDict 中每个元素的 log2() 值。

log2_() → T¶

就地计算 TensorDict 中每个元素的 log2() 值。

log_() → T¶

就地计算 TensorDict 中每个元素的 log() 值。

abstract make_memmap(key: NestedKey, shape: torch.Size | torch.Tensor, *, dtype: torch.dtype | None = None) → MemoryMappedTensor¶

根据形状和可选的数据类型创建空的内存映射张量。

警告

此方法在设计上不是锁安全的。在多个节点上存在的内存映射 TensorDict 实例需要使用 memmap_refresh_() 方法更新。

写入现有条目会导致错误。

参数:

key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，则会引发异常。

shape (torch.Size 或等效, torch.Tensor 用于嵌套张量) – 要写入的张量的形状。

关键字参数:

dtype (torch.dtype, 可选) – 新张量的 dtype。

返回值：:

新的内存映射张量。

abstract make_memmap_from_storage(key: NestedKey, storage: torch.UntypedStorage, shape: torch.Size | torch.Tensor, *, dtype: torch.dtype | None = None) → MemoryMappedTensor¶

根据存储、形状和可选的数据类型创建空的内存映射张量。

警告

此方法在设计上不是锁安全的。在多个节点上存在的内存映射 TensorDict 实例需要使用 memmap_refresh_() 方法更新。

注意

如果存储关联了文件名，则该文件名必须与文件的新的文件名匹配。如果没有关联文件名，但 tensordict 关联了路径，这会导致异常。

参数:

key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，则会引发异常。

storage (torch.UntypedStorage) – 用于新的 MemoryMappedTensor 的存储。必须是物理内存存储。

shape (torch.Size 或等效, torch.Tensor 用于嵌套张量) – 要写入的张量的形状。

关键字参数:

dtype (torch.dtype, 可选) – 新张量的 dtype。

返回值：:

使用给定存储的新内存映射张量。

abstract make_memmap_from_tensor(key: NestedKey, tensor: Tensor, *, copy_data: bool = True) → MemoryMappedTensor¶

根据张量创建空的内存映射张量。

警告

此方法在设计上不是锁安全的。在多个节点上存在的内存映射 TensorDict 实例需要使用 memmap_refresh_() 方法更新。

如果 copy_data 为 True（即存储未共享），此方法始终复制存储内容。

参数:

key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，则会引发异常。

tensor (torch.Tensor) – 要在物理内存上复制的张量。

关键字参数:

copy_data (bool, optionaL) – 如果为 False，新的张量将共享输入的元数据，例如形状和 dtype，但内容将为空。默认值为 True。

返回值：:

使用给定存储的新内存映射张量。

map(fn: Callable[[TensorDictBase], TensorDictBase | None], dim: int = 0, num_workers: int | None = None, *, out: TensorDictBase | None = None, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: torch.Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = False, pbar: bool = False, mp_start_method: str | None = None)¶

将函数映射到 tensordict 跨一个维度的拆分。

此方法将通过将 tensordict 实例分成大小相等的 tensordict 并根据所需的工作进程数量分派操作，将函数应用于 tensordict 实例。

函数签名应为 Callabe[[TensorDict], Union[TensorDict, Tensor]]。输出必须支持 torch.cat() 操作。函数必须可序列化。

参数:

fn (callable) – 要应用于 tensordict 的函数。支持类似于 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的签名。

dim (int, optional) – tensordict 将在其上进行分块的维度。

num_workers (int, optional) – 工作进程数量。与 pool 互斥。如果没有提供，工作进程数量将设置为可用 cpu 数量。

关键字参数:

out (TensorDictBase, optional) – 输出的可选容器。它沿提供的 dim 的批次大小必须与 self.ndim 匹配。如果它是共享的或内存映射的（is_shared() 或 is_memmap() 返回 True），它将在远程进程中填充，避免数据向内传输。否则，来自 self 切片的数据将被发送到进程，在当前进程中收集，并原地写入到 out 中。

chunksize (int, optional) – 每个数据块的大小。chunksize 为 0 将沿所需维度解绑定 tensordict，并在应用函数后将其重新堆叠，而 chunksize>0 将拆分 tensordict 并对生成的 tensordict 列表调用 torch.cat()。如果没有提供，块数将等于工作进程数量。对于非常大的 tensordict，如此大的块可能无法在内存中容纳以完成操作，可能需要更多块才能使操作在实践中可行。此参数与 num_chunks 互斥。

num_chunks (int, optional) – 将 tensordict 拆分成多少块。如果没有提供，块数将等于工作进程数量。对于非常大的 tensordict，如此大的块可能无法在内存中容纳以完成操作，可能需要更多块才能使操作在实践中可行。此参数与 chunksize 互斥。

pool (mp.Pool, optional) – 要用来执行作业的多进程池实例。如果没有提供，将在 map 方法中创建池。

generator (torch.Generator, optional) –
要用来种子的生成器。一个基础种子将从它生成，并且池的每个工作进程都将使用提供的种子加上一个从 0 到 num_workers 的唯一整数进行播种。如果没有提供生成器，将使用一个随机整数作为种子。要使用未播种的工作进程，应该单独创建池并直接传递给 map()。.. note

Caution should be taken when providing a low-valued seed as this can cause autocorrelation between experiments, example: if 8 workers are asked and the seed is 4, the workers seed will range from 4 to 11. If the seed is 5, the workers seed will range from 5 to 12. These two experiments will have an overlap of 7 seeds, which can have unexpected effects on the results.

注意

工作进程播种的目的是让每个工作进程都有独立的种子，而不是在调用 map 方法时获得可重复的结果。换句话说，两个实验可能会（并且很可能）返回不同的结果，因为不可能知道哪个工作进程将选择哪个作业。但是，我们可以确保每个工作进程都有不同的种子，并且每个工作进程上的伪随机操作将是不相关的。

max_tasks_per_child (int, optional) – 每个子进程选取的最大作业数。默认值为 None，即对作业数没有限制。

worker_threads (int, optional) – 工作线程的数量。默认值为 1。

index_with_generator (bool, optional) – 如果为 True，则在查询期间进行 TensorDict 的拆分/分块，以节省初始化时间。请注意，chunk() 和 split() 比索引（在生成器中使用）效率高得多，因此初始化时间的处理时间收益可能会对总运行时间产生负面影响。默认值为 False。

pbar (bool, optional) – 如果为 True，则会显示进度条。需要 tqdm 可用。默认值为 False。

mp_start_method (str, optional) – 多进程的启动方法。如果没有提供，将使用默认启动方法。可接受的字符串为 "fork" 和 "spawn"。请记住，使用 "fork" 启动方法，进程之间无法共享 "cuda" 张量。如果 pool 传递给 map 方法，则此方法无效。

示例

>>> import torch >>> from tensordict import TensorDict >>> >>> def process_data(data): ... data.set("y", data.get("x") + 1) ... return data >>> if __name__ == "__main__": ... data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_() ... data = data.map(process_data, dim=1) ... print(data["y"][:, :10]) ... tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

注意

此方法在处理存储在磁盘上的大型数据集（例如内存映射 TensorDict）时特别有用，其中块将是原始数据的零复制切片，这些切片几乎可以零成本传递给进程。这允许以很小的成本处理非常大的数据集（例如超过 Tb 大的数据集）。

map_iter(fn: Callable[[TensorDictBase], TensorDictBase | None], dim: int = 0, num_workers: int | None = None, *, shuffle: bool = False, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: torch.Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = True, pbar: bool = False, mp_start_method: str | None = None)¶

将一个函数映射到 TensorDict 的一个维度上的拆分，并进行迭代。

这是 map() 的可迭代版本。

此方法通过将 TensorDict 实例切分成大小相等的 TensorDict，并在所需数量的工作线程上调度操作，将函数应用于 TensorDict 实例。它将一次返回一个结果。

函数签名应为 Callabe[[TensorDict], Union[TensorDict, Tensor]]。该函数必须是可序列化的。

参数:

fn (callable) – 要应用于 tensordict 的函数。支持类似于 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的签名。

dim (int, optional) – tensordict 将在其上进行分块的维度。

num_workers (int, optional) – 工作进程数量。与 pool 互斥。如果没有提供，工作进程数量将设置为可用 cpu 数量。

关键字参数:

shuffle (bool, optional) – 是否全局打乱索引。如果为 True，则每个批次将包含非连续样本。如果 index_with_generator=False 并且 shuffle=True`，则会引发错误。默认值为 False。

chunksize (int, optional) – 每个数据块的大小。chunksize 为 0 将沿所需维度解绑定 tensordict，并在应用函数后将其重新堆叠，而 chunksize>0 将拆分 tensordict 并对生成的 tensordict 列表调用 torch.cat()。如果没有提供，块数将等于工作进程数量。对于非常大的 tensordict，如此大的块可能无法在内存中容纳以完成操作，可能需要更多块才能使操作在实践中可行。此参数与 num_chunks 互斥。

num_chunks (int, optional) – 将 tensordict 拆分成多少块。如果没有提供，块数将等于工作进程数量。对于非常大的 tensordict，如此大的块可能无法在内存中容纳以完成操作，可能需要更多块才能使操作在实践中可行。此参数与 chunksize 互斥。

pool (mp.Pool, optional) – 要用来执行作业的多进程池实例。如果没有提供，将在 map 方法中创建池。

generator (torch.Generator, optional) –
要用来种子的生成器。一个基础种子将从它生成，并且池的每个工作进程都将使用提供的种子加上一个从 0 到 num_workers 的唯一整数进行播种。如果没有提供生成器，将使用一个随机整数作为种子。要使用未播种的工作进程，应该单独创建池并直接传递给 map()。.. note

Caution should be taken when providing a low-valued seed as this can cause autocorrelation between experiments, example: if 8 workers are asked and the seed is 4, the workers seed will range from 4 to 11. If the seed is 5, the workers seed will range from 5 to 12. These two experiments will have an overlap of 7 seeds, which can have unexpected effects on the results.

注意

工作进程播种的目的是让每个工作进程都有独立的种子，而不是在调用 map 方法时获得可重复的结果。换句话说，两个实验可能会（并且很可能）返回不同的结果，因为不可能知道哪个工作进程将选择哪个作业。但是，我们可以确保每个工作进程都有不同的种子，并且每个工作进程上的伪随机操作将是不相关的。

max_tasks_per_child (int, optional) – 每个子进程选取的最大作业数。默认值为 None，即对作业数没有限制。

worker_threads (int, optional) – 工作线程的数量。默认值为 1。

index_with_generator (bool, optional) –
如果为 True，则在查询期间进行 TensorDict 的拆分/分块，以节省初始化时间。请注意，chunk() 和 split() 比索引（在生成器中使用）效率高得多，因此初始化时间的处理时间收益可能会对总运行时间产生负面影响。默认值为 True。

注意

index_with_generator 的默认值对于 map_iter 和 map 不同，前者假设在内存中存储 TensorDict 的拆分版本成本过高。

pbar (bool, optional) – 如果为 True，则会显示进度条。需要 tqdm 可用。默认值为 False。

mp_start_method (str, optional) – 多进程的启动方法。如果没有提供，将使用默认启动方法。可接受的字符串为 "fork" 和 "spawn"。请记住，使用 "fork" 启动方法，进程之间无法共享 "cuda" 张量。如果 pool 传递给 map 方法，则此方法无效。

示例

>>> import torch >>> from tensordict import TensorDict >>> >>> def process_data(data): ... data.unlock_() ... data.set("y", data.get("x") + 1) ... return data >>> if __name__ == "__main__": ... data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_() ... for sample in data.map_iter(process_data, dim=1, chunksize=5): ... print(sample["y"]) ... break ... tensor([[1., 1., 1., 1., 1.]])

注意

此方法在处理存储在磁盘上的大型数据集（例如内存映射 TensorDict）时特别有用，其中块将是原始数据的零复制切片，这些切片几乎可以零成本传递给进程。这允许以很小的成本处理非常大的数据集（例如超过 Tb 大的数据集）。

注意

此函数可用于表示数据集并以类似于数据加载器的方式从中加载数据。

abstract masked_fill(mask: Tensor, value: float | bool) → T¶

masked_fill 的非就地版本。

参数:

mask (boolean torch.Tensor) – 要填充的值的掩码。形状必须与 TensorDict 批次大小匹配。

value – 用于填充张量的值。

返回值：:

self

示例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)}, ... batch_size=[3]) >>> mask = torch.tensor([True, False, False]) >>> td1 = td.masked_fill(mask, 1.0) >>> td1.get("a") tensor([[1., 1., 1., 1.], [0., 0., 0., 0.], [0., 0., 0., 0.]])

abstract masked_fill_(mask: Tensor, value: float | bool) → T¶

用指定的值填充与掩码相对应的值。

参数:

mask (boolean torch.Tensor) – 要填充的值的掩码。形状必须与 TensorDict 批次大小匹配。

value – 用于填充张量的值。

返回值：:

self

示例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)}, ... batch_size=[3]) >>> mask = torch.tensor([True, False, False]) >>> td.masked_fill_(mask, 1.0) >>> td.get("a") tensor([[1., 1., 1., 1.], [0., 0., 0., 0.], [0., 0., 0., 0.]])

abstract masked_select(mask: Tensor) → T¶

对 TensorDict 中的所有张量进行掩码，并返回一个新的 TensorDict 实例，其中类似的键指向掩码后的值。

参数:

mask (torch.Tensor) – 用于张量的布尔掩码。形状必须与 TensorDict 的 batch_size 相匹配。

示例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)}, ... batch_size=[3]) >>> mask = torch.tensor([True, False, False]) >>> td_mask = td.masked_select(mask) >>> td_mask.get("a") tensor([[0., 0., 0., 0.]])

maximum(other: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T¶

计算 self 和 other 的逐元素最大值。

参数:

other (TensorDict or Tensor) – 另一个输入 tensordict 或张量。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

maximum_(other: TensorDictBase | torch.Tensor) → T¶

maximum() 的就地版本。

注意

就地 maximum 不支持 default 关键字参数。

classmethod maybe_dense_stack(input, dim=0, *, out=None, **kwargs)¶

尝试对 tensordicts 进行密集堆叠，并在需要时回退到惰性堆叠。

有关详细信息，请参阅 maybe_dense_stack()。

mean(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: torch.dtype | None = None, reduce: bool | None = None) → TensorDictBase | torch.Tensor¶

返回输入 tensordict 中所有元素的平均值。

参数:

dim (int, tuple of int, optional) – 如果为 None，则返回一个无维度的 tensordict，其中包含所有叶子的平均值（如果可以计算）。如果为整数或整数元组，则只有当此维度与 tensordict 形状兼容时，才会在指定的维度上调用 mean。

keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上进行缩减，并返回一个单个缩减后的张量。默认为 False。

memmap(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

将所有张量写入一个新的 tensordict 中与之对应的内存映射张量。

参数:

prefix (str) – 内存映射张量将存储在其中的目录前缀。目录树结构将模仿 tensordict 的结构。

copy_existing (bool) – 如果为 False（默认值），则如果 tensordict 中的条目已经是磁盘上存储的张量，并且具有关联的文件，但未根据前缀保存在正确的位置，则会引发异常。如果为 True，则任何现有的张量都将被复制到新位置。

关键字参数:

num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0。

return_early (bool, optional) – 如果为 True 且 num_threads>0，则该方法将返回 tensordict 的 future。

share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程之间共享，并且对单个节点内任何工作程序的写入操作（例如就地更新或设置）将更新所有其他工作程序上的值。如果非张量叶节点数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。

然后 TensorDict 被锁定，这意味着任何非就地写入操作都会抛出异常（例如，重命名、设置或删除条目）。一旦 tensordict 解锁，内存映射属性将变为 False，因为不再保证跨进程身份。

返回值：:

如果 return_early=False，则使用存储在磁盘上的张量创建一个新的 tensordict，否则为一个 TensorDictFuture 实例。

注意

以这种方式序列化可能在深度嵌套的 tensordict 中很慢，因此建议不要在训练循环中调用此方法。

memmap_(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

将所有张量就地写入相应的内存映射张量。

参数:

prefix (str) – 内存映射张量将存储在其中的目录前缀。目录树结构将模仿 tensordict 的结构。

copy_existing (bool) – 如果为 False（默认值），则如果 tensordict 中的条目已经是磁盘上存储的张量，并且具有关联的文件，但未根据前缀保存在正确的位置，则会引发异常。如果为 True，则任何现有的张量都将被复制到新位置。

关键字参数:

num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0。

return_early (bool, optional) – 如果为 True 且 num_threads>0，则该方法将返回 tensordict 的 future。可以使用 future.result() 查询结果 tensordict。

share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程之间共享，并且对单个节点内任何工作程序的写入操作（例如就地更新或设置）将更新所有其他工作程序上的值。如果非张量叶节点数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。

然后 TensorDict 被锁定，这意味着任何非就地写入操作都会抛出异常（例如，重命名、设置或删除条目）。一旦 tensordict 解锁，内存映射属性将变为 False，因为不再保证跨进程身份。

返回值：:

如果 return_early=False，则为 self，否则为 TensorDictFuture 实例。

注意

以这种方式序列化可能在深度嵌套的 tensordict 中很慢，因此建议不要在训练循环中调用此方法。

memmap_like(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

创建与原始 tensordict 形状相同的无内容内存映射 tensordict。

参数:

prefix (str) – 内存映射张量将存储在其中的目录前缀。目录树结构将模仿 tensordict 的结构。

copy_existing (bool) – 如果为 False（默认值），则如果 tensordict 中的条目已经是磁盘上存储的张量，并且具有关联的文件，但未根据前缀保存在正确的位置，则会引发异常。如果为 True，则任何现有的张量都将被复制到新位置。

关键字参数:

num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0。

return_early (bool, optional) – 如果为 True 且 num_threads>0，则该方法将返回 tensordict 的 future。

share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程之间共享，并且对单个节点内任何工作程序的写入操作（例如就地更新或设置）将更新所有其他工作程序上的值。如果非张量叶节点数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。

然后 TensorDict 被锁定，这意味着任何非就地写入操作都会抛出异常（例如，重命名、设置或删除条目）。一旦 tensordict 解锁，内存映射属性将变为 False，因为不再保证跨进程身份。

返回值：:

如果 return_early=False，则为一个新的 TensorDict 实例，其数据存储为内存映射张量，否则为一个 TensorDictFuture 实例。

注意

这是在磁盘上写入一组大型缓冲区的推荐方法，因为 memmap_() 会复制信息，这对于大型内容来说可能很慢。

示例

>>> td = TensorDict({ ... "a": torch.zeros((3, 64, 64), dtype=torch.uint8), ... "b": torch.zeros(1, dtype=torch.int64), ... }, batch_size=[]).expand(1_000_000) # expand does not allocate new memory >>> buffer = td.memmap_like("/path/to/dataset")

memmap_refresh_()¶

如果内存映射 tensordict 具有 saved_path，则刷新其内容。

如果与其关联的路径不存在，此方法将引发异常。

minimum(other: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T¶

计算 self 和 other 的逐元素最小值。

参数:

other (TensorDict or Tensor) – 另一个输入 tensordict 或张量。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

minimum_(other: TensorDictBase | torch.Tensor) → T¶

minimum() 的就地版本。

注意

就地 minimum 不支持 default 关键字参数。

mul(other: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T¶

将 other 乘以 self。

\[\text{{out}}_i = \text{{input}}_i \times \text{{other}}_i\]

支持广播、类型提升以及整数、浮点和复数输入。

参数:

other (TensorDict, Tensor 或 Number) – 从 self 中减去的张量或数字。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

mul_(other: TensorDictBase | torch.Tensor) → T¶

mul() 的就地版本。

注意

就地 mul 不支持 default 关键字参数。

named_apply(fn: Callable, *others: T, nested_keys: bool = False, batch_size: Sequence[int] | None = None, device: torch.device | None = _NoDefault.ZERO, names: Sequence[str] | None = _NoDefault.ZERO, inplace: bool = False, default: Any = _NoDefault.ZERO, filter_empty: bool | None = None, propagate_lock: bool = False, call_on_nested: bool = False, out: TensorDictBase | None = None, **constructor_kwargs) → T | None¶

将一个键值条件的调用应用于 tensordict 中存储的所有值，并将它们设置在一个新的 atensordict 中。

可调用签名的格式必须为 Callable[Tuple[str, Tensor, ...], Optional[Union[Tensor, TensorDictBase]]].

参数:

fn (Callable) – 要应用于 tensordict 中 (name, tensor) 对的函数。对于每个叶子，只使用它的叶子名称（而不是完整的 NestedKey）。

*others (TensorDictBase 实例, 可选) – 如果提供，这些 tensordict 实例应该具有与 self 相匹配的结构。 fn 参数应该接收与 tensordict 数量相同的未命名输入，包括 self。如果其他 tensordict 缺少条目，可以通过 default 关键字参数传递默认值。

nested_keys (bool, optional) – 如果为 True，则使用到叶子的完整路径。默认为 False，即只将最后一个字符串传递给函数。

batch_size (int 序列, 可选) – 如果提供，则生成的 TensorDict 将具有所需的 batch_size。 batch_size 参数应该与转换后的 batch_size 相匹配。这是一个关键字参数。

device (torch.device, 可选) – 生成结果的设备，如果有。

names (字符串列表, 可选) – 新的维度名称，如果修改了 batch_size。

inplace (bool, optional) – 如果为 True，则进行就地修改。默认值为 False。这是一个关键字 only 参数。

default (任何类型, 可选) – 其他 tensordict 中缺少条目的默认值。如果未提供，则缺少的条目将引发 KeyError。

filter_empty (bool, optional) – 如果为 True，则会过滤掉空的 tensordict。这同时也带来了更低的计算成本，因为不会创建和销毁空数据结构。为了向后兼容，默认为 False。

propagate_lock (bool, optional) – 如果为 True，则锁定 tensordict 会产生另一个锁定 tensordict。默认为 False。

call_on_nested (bool, optional) –
如果 True，则该函数将在第一级张量和容器（TensorDict 或张量类）上调用。在这种情况下，func 负责将调用传播到嵌套级别。这允许在将调用传播到嵌套 tensordict 时实现细粒度行为。如果 False，则该函数仅在叶节点上调用，并且 apply 将负责将该函数分派到所有叶节点。

>>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]}) >>> def mean_tensor_only(val): ... if is_tensor_collection(val): ... raise RuntimeError("Unexpected!") ... return val.mean() >>> td_mean = td.apply(mean_tensor_only) >>> def mean_any(val): ... if is_tensor_collection(val): ... # Recurse ... return val.apply(mean_any, call_on_nested=True) ... return val.mean() >>> td_mean = td.apply(mean_any, call_on_nested=True)

out (TensorDictBase, 可选) –
一个 tensordict，用于写入结果。这可以用来避免创建新的 tensordict

>>> td = TensorDict({"a": 0}) >>> td.apply(lambda x: x+1, out=td) >>> assert (td==1).all()

警告

如果 tensordict 上执行的操作需要访问多个键才能进行单个计算，则提供一个等于 self 的 out 参数会导致该操作静默地提供错误的结果。例如

>>> td = TensorDict({"a": 1, "b": 1}) >>> td.apply(lambda x: x+td["a"])["b"] # Right! tensor(2) >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong! tensor(3)

**constructor_kwargs** – 传递给 TensorDict 构造函数的额外关键字参数。

返回值：:

一个新的 tensordict，包含已转换的张量。

示例

>>> td = TensorDict({ ... "a": -torch.ones(3), ... "nested": {"a": torch.ones(3), "b": torch.zeros(3)}}, ... batch_size=[3]) >>> def name_filter(name, tensor): ... if name == "a": ... return tensor >>> td.named_apply(name_filter) TensorDict( fields={ a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), nested: TensorDict( fields={ a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False) >>> def name_filter(name, *tensors): ... if name == "a": ... r = 0 ... for tensor in tensors: ... r = r + tensor ... return tensor >>> out = td.named_apply(name_filter, td) >>> print(out) TensorDict( fields={ a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), nested: TensorDict( fields={ a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False)}, batch_size=torch.Size([3]), device=None, is_shared=False) >>> print(out["a"]) tensor([-1., -1., -1.])

注意

如果函数返回 None，则忽略该条目。这可以用来过滤 tensordict 中的数据

>>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, []) >>> def name_filter(name, tensor): ... if name == "1": ... return tensor >>> td.named_apply(name_filter) TensorDict( fields={ 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)

abstract property names¶

tensordict 的维度名称。

可以使用 names 参数在构造时设置名称。

有关在构造后如何设置名称的详细信息，请参阅 refine_names()。

nanmean(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: torch.dtype | None = None, reduce: bool | None = None) → TensorDictBase | torch.Tensor¶

返回输入 tensordict 中所有非 NaN 元素的平均值。

参数:

dim (int, tuple of int, optional) – 如果为 None，则返回一个无维度的 tensordict，其中包含所有叶子的平均值（如果可以计算）。如果为整数或整数元组，则只有当此维度与 tensordict 形状兼容时，才会在指定的维度上调用 mean。

keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上进行缩减，并返回一个单个缩减后的张量。默认为 False。

nansum(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: torch.dtype | None = None, reduce: bool | None = None) → TensorDictBase | torch.Tensor¶

返回输入张量字典中所有非 NaN 元素的总和。

参数:

dim (int, tuple of int, 可选) – 如果为 None，则返回一个无维度的张量字典，其中包含所有叶子节点的总和值（如果可以计算）。如果为整数或整数元组，则当且仅当此维度与张量字典形状兼容时，才会在指定维度上调用 sum。

keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上进行缩减，并返回一个单个缩减后的张量。默认为 False。

property ndim: int¶

参见 batch_dims().

ndimension() → int¶

参见 batch_dims().

neg() → T¶

计算张量字典中每个元素的 neg() 值。

neg_() → T¶

就地计算张量字典中每个元素的 neg() 值。

new_empty(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶

返回一个大小为 size 且包含空张量的张量字典。

默认情况下，返回的张量字典具有与该张量字典相同的 torch.dtype 和 torch.device。

参数:

size (int...) – 定义输出张量形状的整数列表、元组或 torch.Size。

关键字参数:

dtype (torch.dtype, 可选) – 返回张量字典的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

device (torch.device, 可选) – 返回张量字典的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

requires_grad (bool, 可选) – 是否应记录对返回张量的操作的自动梯度。默认值：False。

layout (torch.layout, 可选) – 返回张量字典值的所需布局。默认值：torch.strided。

pin_memory (bool, 可选) – 如果设置，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。

new_full(size: Size, fill_value, *, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶

返回一个大小为 size 且填充为 1 的张量字典。

默认情况下，返回的张量字典具有与该张量字典相同的 torch.dtype 和 torch.device。

参数:

size (整数序列) – 定义输出张量形状的整数列表、元组或 torch.Size。

fill_value (标量) – 用于填充输出张量的数字。

关键字参数:

dtype (torch.dtype, 可选) – 返回张量字典的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

device (torch.device, 可选) – 返回张量字典的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

requires_grad (bool, 可选) – 是否应记录对返回张量的操作的自动梯度。默认值：False。

layout (torch.layout, 可选) – 返回张量字典值的所需布局。默认值：torch.strided。

pin_memory (bool, 可选) – 如果设置，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。

new_ones(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶

返回一个大小为 size 且填充为 1 的张量字典。

默认情况下，返回的张量字典具有与该张量字典相同的 torch.dtype 和 torch.device。

参数:

size (int...) – 定义输出张量形状的整数列表、元组或 torch.Size。

关键字参数:

dtype (torch.dtype, 可选) – 返回张量字典的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

device (torch.device, 可选) – 返回张量字典的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

requires_grad (bool, 可选) – 是否应记录对返回张量的操作的自动梯度。默认值：False。

layout (torch.layout, 可选) – 返回张量字典值的所需布局。默认值：torch.strided。

pin_memory (bool, 可选) – 如果设置，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。

new_tensor(data: torch.Tensor | TensorDictBase, *, dtype: torch.dtype = None, device: DeviceType = _NoDefault.ZERO, requires_grad: bool = False, pin_memory: bool | None = None)¶

返回一个新的 TensorDict，其中 data 为张量 data。

默认情况下，返回的 TensorDict 值具有与该张量相同的 torch.dtype 和 torch.device。

data 也可以是张量集合 (TensorDict 或 tensorclass)，在这种情况下，new_tensor 方法遍历 self 和 data 的张量对。

参数:

data (torch.Tensor 或 TensorDictBase) – 要复制的数据。

关键字参数:

dtype (torch.dtype, 可选) – 返回张量字典的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

device (torch.device, 可选) – 返回张量字典的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

requires_grad (bool, 可选) – 是否应记录对返回张量的操作的自动梯度。默认值：False。

pin_memory (bool, 可选) – 如果设置，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。

new_zeros(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶

返回一个大小为 size 并填充为 0 的 TensorDict。

默认情况下，返回的张量字典具有与该张量字典相同的 torch.dtype 和 torch.device。

参数:

size (int...) – 定义输出张量形状的整数列表、元组或 torch.Size。

关键字参数:

dtype (torch.dtype, 可选) – 返回张量字典的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

device (torch.device, 可选) – 返回张量字典的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

requires_grad (bool, 可选) – 是否应记录对返回张量的操作的自动梯度。默认值：False。

layout (torch.layout, 可选) – 返回张量字典值的所需布局。默认值：torch.strided。

pin_memory (bool, 可选) – 如果设置，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。

non_tensor_items(include_nested: bool = False)¶

返回所有非张量叶子节点，可能递归地。

norm(*, out=None, dtype: torch.dtype | None = None)¶

计算 tensordict 中每个张量的范数。

关键字参数:

out (TensorDict, 可选) – 输出 tensordict。

dtype (torch.dtype, 可选) – 输出 dtype (torch>=2.4)。

numel() → int¶

批次中的元素总数。

下限为 1，因为两个空形状的 tensordict 的堆栈将有两个元素，因此我们认为 tensordict 至少包含 1 个元素。

numpy()¶

将 tensordict 转换为 (可能是嵌套的) numpy 数组字典。

非张量数据按原样公开。

示例

>>> from tensordict import TensorDict >>> import torch >>> data = TensorDict({"a": {"b": torch.zeros(()), "c": "a string!"}}) >>> print(data) TensorDict( fields={ a: TensorDict( fields={ b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), c: NonTensorData(data=a string!, batch_size=torch.Size([]), device=None)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> print(data.numpy()) {'a': {'b': array(0., dtype=float32), 'c': 'a string!'}}

permute(*dims: int)¶

permute(dims: list | tuple)

返回一个 tensordict 的视图，其中批处理维度根据 dims 进行置换。

参数:

*dims_list (int) – tensordict 批处理维度的新的排序方式。或者，可以提供一个包含整数的单个可迭代对象。

dims (list of int) – 调用 permute(…) 的另一种方式。

返回值：:

一个新的 tensordict，其批处理维度按所需顺序排列。

示例

>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4]) >>> print(tensordict.permute([1, 0])) PermutedTensorDict( source=TensorDict( fields={ a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)}, batch_size=torch.Size([3, 4]), device=cpu, is_shared=False), op=permute(dims=[1, 0])) >>> print(tensordict.permute(1, 0)) PermutedTensorDict( source=TensorDict( fields={ a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)}, batch_size=torch.Size([3, 4]), device=cpu, is_shared=False), op=permute(dims=[1, 0])) >>> print(tensordict.permute(dims=[1, 0])) PermutedTensorDict( source=TensorDict( fields={ a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)}, batch_size=torch.Size([3, 4]), device=cpu, is_shared=False), op=permute(dims=[1, 0]))

pin_memory(num_threads: int | None = None, inplace: bool = False) → T¶

在存储的张量上调用 pin_memory()。

参数:

num_threads (int or str) – 如果提供，则在叶子节点上调用 pin_memory 时要使用的线程数。默认为 None，这将在 ThreadPoolExecutor(max_workers=None) 中设置一个较高的线程数。要在线程上执行对 pin_memory() 的所有调用，请传递 num_threads=0。

inplace (bool, optional) – 如果为 True，则 tensordict 将被就地修改。默认为 False。

pin_memory_(num_threads: int | str = 0) → T¶

在存储的张量上调用 pin_memory() 并返回就地修改的 TensorDict。

参数:

num_threads (int or str) – 如果提供，则在叶子节点上调用 pin_memory 时要使用的线程数。如果传递了 "auto"，则将自动确定线程数。

pop(key: NestedKey, default: Any = _NoDefault.ZERO) → Tensor¶

从 tensordict 中删除并返回一个值。

如果值不存在且未提供默认值，则会抛出 KeyError。

参数:

key (str or nested key) – 要查找的条目。

default (Any, optional) – 如果找不到键，则要返回的值。

示例

>>> td = TensorDict({"1": 1}, []) >>> one = td.pop("1") >>> assert one == 1 >>> none = td.pop("1", default=None) >>> assert none is None

abstract popitem() → Tuple[NestedKey, Tensor]¶

删除最后插入 TensorDict 的项目。

popitem 只会返回非嵌套值。

pow(other: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T¶

用 other 对 self 中的每个元素进行幂运算，并返回一个包含结果的张量。

other 可以是单个 float 数、一个 Tensor 或一个 TensorDict。

当 other 是一个张量时，input 和 other 的形状必须是可广播的。

参数:

other (float, tensor or tensordict) – 指数值

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

pow_(other: TensorDictBase | torch.Tensor) → T¶

pow() 的就地版本。

注意

就地 pow 不支持 default 关键字参数。

prod(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: torch.dtype | None = None, reduce: bool | None = None) → TensorDictBase | torch.Tensor¶

返回输入张量字典中所有元素值的乘积。

参数:

dim (int, int 元组, 可选) – 如果为 None，则返回一个无维度的张量字典，其中包含所有叶节点的乘积值（如果可以计算）。如果为整数或整数元组，则仅当该维度与张量字典形状兼容时，才对指定的维度调用 prod。

keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上进行缩减，并返回一个单个缩减后的张量。默认为 False。

qint32()¶

将所有张量转换为 torch.qint32。

qint8()¶

将所有张量转换为 torch.qint8。

quint4x2()¶

将所有张量转换为 torch.quint4x2。

quint8()¶

将所有张量转换为 torch.quint8。

reciprocal() → T¶

计算 TensorDict 中每个元素的 reciprocal() 值。

reciprocal_() → T¶

计算 TensorDict 中每个元素的 reciprocal() 值，并进行就地操作。

recv(src: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → int¶

接收张量字典的内容并用它更新内容。

有关上下文，请查看 send 方法中的示例。

参数:

src (int) – 源工作者的等级。

关键字参数:

group (torch.distributed.ProcessGroup, optional) – 如果设置，则使用指定的进程组进行通信。否则，将使用默认进程组。默认为 None。

init_tag (int) – 源工作者使用的 init_tag。

pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许从不同的节点发送多个数据而不会重叠。请注意，这些伪随机数的生成很昂贵（1e-5 秒/个数），这意味着它可能会减慢算法的运行时间。此值必须与传递给 send() 的值匹配。默认为 False。

reduce(dst, op=None, async_op=False, return_premature=False, group=None)¶

在所有机器上减少张量字典。

只有 rank 为 dst 的进程将接收最终结果。

refine_names(*names) → T¶

根据 names 细化 self 的维度名称。

细化是重命名的特例，它“提升”未命名的维度。None 维度可以细化为任何名称；已命名维度只能细化为具有相同名称的维度。

因为命名张量可以与未命名张量共存，所以细化名称提供了一种编写与命名张量和未命名张量都兼容的命名张量感知代码的好方法。

names 可以包含最多一个 Ellipsis (…）。Ellipsis 会贪婪地展开；它会就地展开以使用来自 self.names 的对应索引的名称来填充与 self.dim() 相同长度的 names。

返回：具有根据输入命名的维度的相同张量字典。

示例

>>> td = TensorDict({}, batch_size=[3, 4, 5, 6]) >>> tdr = td.refine_names(None, None, None, "d") >>> assert tdr.names == [None, None, None, "d"] >>> tdr = td.refine_names("a", None, None, "d") >>> assert tdr.names == ["a", None, None, "d"]

rename(*names, **rename_map)¶

返回一个具有重命名维度的张量字典的克隆。

示例

>>> td = TensorDict({}, batch_size=[1, 2, 3 ,4]) >>> td.names = list("abcd") >>> td_rename = td.rename(c="g") >>> assert td_rename.names == list("abgd")

rename_(*names, **rename_map)¶

与 rename() 相同，但会就地执行重命名。

示例

>>> td = TensorDict({}, batch_size=[1, 2, 3 ,4]) >>> td.names = list("abcd") >>> assert td.rename_(c="g") >>> assert td.names == list("abgd")

abstract rename_key_(old_key: NestedKey, new_key: NestedKey, safe: bool = False) → T¶

使用新的字符串重命名键，并返回具有更新的键名称的相同张量字典。

参数:

old_key (str 或 嵌套键) – 要重命名的键。

new_key (str 或 嵌套键) – 条目的新名称。

safe (bool, 可选) – 如果为 True，则当新键已存在于 TensorDict 中时会抛出错误。

返回值：:

self

replace(*args, **kwargs)¶

创建张量字典的浅拷贝，其中条目已被替换。

接受一个未命名的参数，该参数必须是 TensorDictBase 子类的字典。此外，可以使用命名关键字参数更新第一级条目。

返回值：:

如果输入非空，则为 self 的副本，其中包含更新的条目。如果提供空字典或不提供字典，并且 kwargs 为空，则返回 self。

requires_grad_(requires_grad=True) → T¶

更改是否应记录此张量上的操作的自动梯度：就地设置此张量的 requires_grad 属性。

返回此张量字典。

参数:

requires_grad (bool, 可选) – 自动梯度是否应记录此张量字典上的操作。默认为 True。

abstract reshape(*shape: int)¶

abstract reshape(shape: list | tuple)

返回一个具有所需形状的连续重塑张量。

参数:

*shape (int) – resulting tensordict 的新形状。

返回值：:

一个具有重塑键的 TensorDict

示例

>>> td = TensorDict({ ... 'x': torch.arange(12).reshape(3, 4), ... }, batch_size=[3, 4]) >>> td = td.reshape(12) >>> print(td['x']) torch.Tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])

round() → T¶

计算 TensorDict 中每个元素的 round() 值。

round_() → T¶

就地计算 TensorDict 中每个元素的 round() 值。

save(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

将张量字典保存到磁盘。

此函数是 memmap() 的代理。

property saved_path¶

返回保存的 memmap TensorDict 存储的路径。

当 is_memmap() 返回 False 时（例如，当 tensordict 被解锁时），此参数无效。

select(*keys: NestedKey, inplace: bool = False, strict: bool = True) → T¶

选择 tensordict 的键，并返回一个仅包含所选键的新 tensordict。

不会复制值：对原始张量字典或新张量字典中任何一个张量的原地修改都会导致两个张量字典都发生变化。

参数:

*keys (str) – 要选择的键

inplace (bool) – 如果为 True，则原地修剪张量字典。默认值为 False。

strict (bool, 可选) – 选择不存在的键时是否返回错误。默认值：True。

返回值：:

一个包含所选键的新 tensordict（如果 inplace=True，则为同一个 tensordict）。

注意

要选择 tensordict 中的键并返回一个没有这些键的 tensordict 版本，请参见 split_keys() 方法。

示例

>>> from tensordict import TensorDict >>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, []) >>> td.select("a", ("b", "c")) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.select("a", "b") TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.select("this key does not exist", strict=False) TensorDict( fields={ }, batch_size=torch.Size([]), device=None, is_shared=False)

send(dst: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → None¶

将 tensordict 的内容发送到远程 worker。

参数:

dst (int) – 目标工作者的等级，内容应发送到该工作者。

关键字参数:

group (torch.distributed.ProcessGroup, optional) – 如果设置，则使用指定的进程组进行通信。否则，将使用默认进程组。默认为 None。

init_tag (int) – 用于标记张量的初始标签。请注意，此标签将根据 TensorDict 中包含的张量数量进行递增。

pseudo_rand (bool) – 如果为 True，标签序列将是伪随机的，允许从不同的节点发送多个数据而不会重叠。请注意，这些伪随机数的生成很昂贵（1e-5 秒/数），这意味着它可能会降低算法的运行时间。默认为 False。

示例

>>> from torch import multiprocessing as mp >>> from tensordict import TensorDict >>> import torch >>> >>> >>> def client(): ... torch.distributed.init_process_group( ... "gloo", ... rank=1, ... world_size=2, ... init_method=f"tcp://127.0.0.1:10003", ... ) ... ... td = TensorDict( ... { ... ("a", "b"): torch.randn(2), ... "c": torch.randn(2, 3), ... "_": torch.ones(2, 1, 5), ... }, ... [2], ... ) ... td.send(0) ... >>> >>> def server(queue): ... torch.distributed.init_process_group( ... "gloo", ... rank=0, ... world_size=2, ... init_method=f"tcp://127.0.0.1:10003", ... ) ... td = TensorDict( ... { ... ("a", "b"): torch.zeros(2), ... "c": torch.zeros(2, 3), ... "_": torch.zeros(2, 1, 5), ... }, ... [2], ... ) ... td.recv(1) ... assert (td != 0).all() ... queue.put("yuppie") ... >>> >>> if __name__=="__main__": ... queue = mp.Queue(1) ... main_worker = mp.Process(target=server, args=(queue,)) ... secondary_worker = mp.Process(target=client) ... ... main_worker.start() ... secondary_worker.start() ... out = queue.get(timeout=10) ... assert out == "yuppie" ... main_worker.join() ... secondary_worker.join()

set(key: NestedKey, item: Tensor, inplace: bool = False, *, non_blocking: bool = False, **kwargs: Any) → T¶

设置一个新的键值对。

参数:

key (str, str 元组) – 要设置的键的名称。

item (torch.Tensor 或等效, TensorDictBase 实例) – 要存储在 tensordict 中的值。

inplace (bool, 可选) – 如果为 True 并且键与 tensordict 中的现有键匹配，则更新将就地发生，针对该键值对。如果 inplace 为 True 并且找不到条目，则会添加它。为了更严格的就地操作，请改用 set_()。默认值：False。

关键字参数:

non_blocking (bool, 可选) – 如果为 True 并且此副本位于不同的设备之间，则副本可能与主机异步发生。

返回值：:

self

示例

>>> td = TensorDict({}, batch_size[3, 4]) >>> td.set("x", torch.randn(3, 4)) >>> y = torch.randn(3, 4, 5) >>> td.set("y", y, inplace=True) # works, even if 'y' is not present yet >>> td.set("y", torch.zeros_like(y), inplace=True) >>> assert (y==0).all() # y values are overwritten >>> td.set("y", torch.ones(5), inplace=True) # raises an exception as shapes mismatch

set_(key: NestedKey, item: Tensor, *, non_blocking: bool = False) → T¶

在保留原始存储空间的情况下，将值设置为现有键。

参数:

key (str) – 值的名称

item (torch.Tensor or 兼容类型, TensorDictBase) – 要存储在 tensordict 中的值

关键字参数:

non_blocking (bool, 可选) – 如果为 True 并且此副本位于不同的设备之间，则副本可能与主机异步发生。

返回值：:

self

示例

>>> td = TensorDict({}, batch_size[3, 4]) >>> x = torch.randn(3, 4) >>> td.set("x", x) >>> td.set_("x", torch.zeros_like(x)) >>> assert (x == 0).all()

set_at_(key: NestedKey, value: Tensor, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], *, non_blocking: bool = False) → T¶

在 index 所示的索引处就地设置值。

参数:

key (str, 字符串元组) – 要修改的键。

value (torch.Tensor) – 要在索引 index 处设置的值

index (int, 张量 or 元组) – 要写入值的索引。

关键字参数:

non_blocking (bool, 可选) – 如果为 True 并且此副本位于不同的设备之间，则副本可能与主机异步发生。

返回值：:

self

示例

>>> td = TensorDict({}, batch_size[3, 4]) >>> x = torch.randn(3, 4) >>> td.set("x", x) >>> td.set_at_("x", value=torch.ones(1, 4), index=slice(1)) >>> assert (x[0] == 1).all()

set_non_tensor(key: NestedKey, value: Any)¶

使用 tensordict.tensorclass.NonTensorData 在 tensordict 中注册非张量值。

可以使用 TensorDictBase.get_non_tensor() 或直接使用 get 来检索该值，这将返回 tensordict.tensorclass.NonTensorData 对象。

返回：self

示例

>>> data = TensorDict({}, batch_size=[]) >>> data.set_non_tensor(("nested", "the string"), "a string!") >>> assert data.get_non_tensor(("nested", "the string")) == "a string!" >>> # regular `get` works but returns a NonTensorData object >>> data.get(("nested", "the string")) NonTensorData( data='a string!', batch_size=torch.Size([]), device=None, is_shared=False)

setdefault(key: NestedKey, default: Tensor, inplace: bool = False) → Tensor¶

如果 key 不在 tensordict 中，则插入 key 条目，其值为 default。

如果 key 在 tensordict 中，则返回 key 的值，否则返回 default。

参数:

key (str or 嵌套键) – 值的名称。

default (torch.Tensor or 兼容类型, TensorDictBase) – 如果该键不存在，则要存储在 tensordict 中的值。

返回值：:

tensordict 中 key 的值。如果之前未设置该键，则将为默认值。

示例

>>> td = TensorDict({}, batch_size=[3, 4]) >>> val = td.setdefault("a", torch.zeros(3, 4)) >>> assert (val == 0).all() >>> val = td.setdefault("a", torch.ones(3, 4)) >>> assert (val == 0).all() # output is still 0

property shape: Size¶

参见 batch_size。

abstract share_memory_() → T¶

将所有张量置于共享内存中。

然后锁定 TensorDict，这意味着任何非就地的写入操作都将引发异常（例如，重命名、设置或删除条目）。相反，一旦 tensordict 解锁，share_memory 属性将变为 False，因为跨进程标识不再保证。

返回值：:

self

sigmoid() → T¶

计算 TensorDict 中每个元素的 sigmoid() 值。

sigmoid_() → T¶

就地计算 TensorDict 中每个元素的 sigmoid() 值。

sign() → T¶

计算 TensorDict 中每个元素的 sign() 值。

sign_() → T¶

就地计算 TensorDict 中每个元素的 sign() 值。

sin() → T¶

计算 TensorDict 中每个元素的 sin() 值。

sin_() → T¶

就地计算 TensorDict 中每个元素的 sin() 值。

sinh() → T¶

计算 TensorDict 中每个元素的 sinh() 值。

sinh_() → T¶

在 TensorDict 中对每个元素的 sinh() 值进行就地计算。

size(dim: int | None = None) → torch.Size | int¶

返回由 dim 指示的维度的尺寸。

如果未指定 dim，则返回 TensorDict 的 batch_size 属性。

property sorted_keys: list[NestedKey]¶

按字母顺序返回排序的键。

不支持额外的参数。

如果 TensorDict 被锁定，则键会被缓存，直到 TensorDict 被解锁以加快执行速度。

abstract split(split_size: int | list[int], dim: int = 0) → list[TensorDictBase]¶

使用给定维度中指定的尺寸，将 TensorDict 中的每个张量分割，类似于 torch.split。

返回一个包含 TensorDict 实例的列表，这些实例是对项目分割块的视图。

参数:

split_size (int 或 List(int)) – 单个块的尺寸或每个块的尺寸列表。

dim (int) – 沿其分割张量的维度。

返回值：:

一个包含在给定维度中具有指定尺寸的 TensorDict 的列表。

示例

>>> td = TensorDict({ ... 'x': torch.arange(12).reshape(3, 4), ... }, batch_size=[3, 4]) >>> td0, td1 = td.split([1, 2], dim=0) >>> print(td0['x']) torch.Tensor([[0, 1, 2, 3]])

split_keys(*key_sets, inplace=False, strict: bool = True, reproduce_struct: bool = False)¶

根据一个或多个键集将 TensorDict 分割成子集。

该方法将返回 N+1 个 TensorDict，其中 N 是提供的参数数量。

参数:

inplace (bool, 可选) – 如果为 True，则键会从 self 中就地删除。默认为 False。

strict (bool, 可选) – 如果为 True，则在键丢失时会引发异常。默认为 True。

reproduce_struct (bool, 可选) – 如果为 True，则返回的所有 TensorDict 都会与 self 具有相同的树结构，即使某些子 TensorDict 不包含叶子。

注意

None 非张量值将被忽略，不会返回。

注意

该方法不会检查提供的列表中的重复项。

示例

>>> td = TensorDict( ... a=0, ... b=0, ... c=0, ... d=0, ... ) >>> td_a, td_bc, td_d = td.split_keys(["a"], ["b", "c"]) >>> print(td_bc)

sqrt()¶

计算 self 的逐元素平方根。

sqrt_()¶

sqrt() 的就地版本。

squeeze(dim: int | None = None) → T¶

压缩 -self.batch_dims+1 到 self.batch_dims-1 之间的维度上的所有张量，并将其返回到一个新的 TensorDict 中。

参数:

dim (可选[int]) – 要压缩的维度。如果 dim 为 None，则所有单例维度都会被压缩。默认为 None。

示例

>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 1, 4, 2), ... }, batch_size=[3, 1, 4]) >>> td = td.squeeze() >>> td.shape torch.Size([3, 4]) >>> td.get("x").shape torch.Size([3, 4, 2])

此操作也可以用作上下文管理器。对原始 TensorDict 的更改将发生在外部，即原始张量的内容不会被更改。这也假定 TensorDict 没有被锁定（否则，需要解锁 TensorDict）。此功能与隐式压缩不兼容。

>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 1, 4, 2), ... }, batch_size=[3, 1, 4]) >>> with td.squeeze(1) as tds: ... tds.set("y", torch.zeros(3, 4)) >>> assert td.get("y").shape == [3, 1, 4]

classmethod stack(input, dim=0, *, out=None)¶

将 TensorDict 沿给定维度堆叠成单个 TensorDict。

此调用等效于调用 torch.stack()，但与 torch.compile 兼容。

state_dict(destination=None, prefix='', keep_vars=False, flatten=False) → OrderedDict[str, Any]¶

从 TensorDict 生成 state_dict。

state-dict 的结构将仍然是嵌套的，除非 flatten 设置为 True。

TensorDict state-dict 包含重建 TensorDict 所需的所有张量和元数据（目前不支持名称）。

参数:

destination (dict, 可选) – 如果提供，则 TensorDict 的状态将更新到字典中，并返回同一个对象。否则，将创建一个 OrderedDict 并返回。默认： None。

prefix (str, 可选) – 添加到张量名称的前缀，以组成 state_dict 中的键。默认： ''。

keep_vars (bool, 可选) – 默认情况下，state dict 中返回的 torch.Tensor 项目与自动微分分离。如果将其设置为 True，则不会执行分离。默认： False。

flatten (bool, 可选) – 结构是否应该使用 "." 字符进行扁平化。默认值为 False。

示例

>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, []) >>> sd = data.state_dict() >>> print(sd) OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3', OrderedDict([('3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])), ('__batch_size', torch.Size([])), ('__device', None)]) >>> sd = data.state_dict(flatten=True) OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3.3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])

std(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: bool | None = None) → TensorDictBase | torch.Tensor¶

返回输入张量字典中所有元素的标准差值。

参数:

dim (int, int 元组, 可选) – 如果为 None，则返回一个无量纲的张量字典，包含所有叶子的求和值（如果可以计算）。如果为整数或整数元组，则只有当此维度与张量字典的形状兼容时，才会在指定的维度上调用 std。

keepdim (bool) – 输出张量是否保留维度。

关键字参数:

correction (int) – 样本大小与样本自由度之间的差值。默认为贝塞尔校正，correction=1。

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上进行缩减，并返回一个单个缩减后的张量。默认为 False。

sub(other: TensorDictBase | float, *, alpha: float | None = None, default: str | CompatibleType | None = None)¶

从 self 中减去 other，并乘以 alpha。

\[\text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i\]

支持广播、类型提升以及整数、浮点和复数输入。

参数:

other (TensorDict, Tensor 或 Number) – 从 self 中减去的张量或数字。

关键字参数:

alpha (数字) – other 的乘数。

default (torch.Tensor 或 str, 可选) – 用于独占条目时的默认值。如果没有提供，两个 tensordicts 的键列表必须完全匹配。如果传递 default="intersection"，则只考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于两侧操作中所有缺失的条目。

sub_(other: TensorDictBase | float, alpha: float | None = None)¶

sub() 的就地版本。

注意

就地 sub 不支持 default 关键字参数。

sum(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: torch.dtype | None = None, reduce: bool | None = None) → TensorDictBase | torch.Tensor¶

返回输入张量字典中所有元素的求和值。

参数:

dim (int, tuple of int, 可选) – 如果为 None，则返回一个无维度的张量字典，其中包含所有叶子节点的总和值（如果可以计算）。如果为整数或整数元组，则当且仅当此维度与张量字典形状兼容时，才会在指定维度上调用 sum。

keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上进行缩减，并返回一个单个缩减后的张量。默认为 False。

tan() → T¶

计算张量字典中每个元素的 tan() 值。

tan_() → T¶

就地计算张量字典中每个元素的 tan() 值。

tanh() → T¶

计算张量字典中每个元素的 tanh() 值。

tanh_() → T¶

就地计算张量字典中每个元素的 tanh() 值。

to(device: Optional[Union[int, device]] = ..., dtype: Optional[Union[device, str]] = ..., non_blocking: bool = ...) → T¶

to(dtype: Union[device, str], non_blocking: bool = ...) → T

to(tensor: Tensor, non_blocking: bool = ...) → T

to(*, other: T, non_blocking: bool = ...) → T

to(*, batch_size: Size) → T

将 TensorDictBase 子类映射到另一个设备、数据类型或另一个 TensorDictBase 子类（如果允许）。

不允许将张量转换为新的数据类型，因为 TensorDict 不绑定到包含单个张量数据类型。

参数:

device (torch.device, optional) – TensorDict 的目标设备。

dtype (torch.dtype, optional) – TensorDict 的目标浮点数或复数数据类型。

tensor (torch.Tensor, optional) – 数据类型和设备是此 TensorDict 中所有张量目标数据类型和设备的张量。

关键字参数:

non_blocking (bool, optional) – 操作是否应阻塞。

memory_format (torch.memory_format, optional) – 此 TensorDict 中 4D 参数和缓冲区的目标内存格式。

batch_size (torch.Size, optional) – 输出 TensorDict 的目标批次大小。

other (TensorDictBase, optional) –
数据类型和设备是此 TensorDict 中所有张量目标数据类型和设备的 TensorDict 实例。.. 注意:: 由于 TensorDictBase 实例没有

数据类型，数据类型是从示例叶节点中收集的。如果有多个数据类型，则不执行数据类型转换。

non_blocking_pin (bool, optional) –
如果 True，则在将张量发送到设备之前将其固定。这将异步完成，但可以通过 num_threads 参数进行控制。

注意

调用 tensordict.pin_memory().to("cuda") 通常比 tensordict.to("cuda", non_blocking_pin=True) 慢得多，因为在第二种情况下异步调用 pin_memory。如果张量很大且数量众多，多线程 pin_memory 通常会有益：当要发送的张量太少时，生成线程和收集数据的开销大于多线程的益处，如果张量很小，遍历长列表的开销也很大。过大。

num_threads (int or None, optional) – 如果 non_blocking_pin=True，则用于 pin_memory 的线程数。默认情况下，将生成 max(1, torch.get_num_threads()) 线程。 num_threads=0 将取消对 pin_memory() 调用的任何多线程。

返回值：:

如果设备不同于 TensorDict 设备，或如果传递了数据类型，则为新的 TensorDict 实例。否则，为相同的 TensorDict。仅对 batch_size 修改进行就地操作。

注意

如果 TensorDict 已合并，则生成的 TensorDict 也将合并。每个新张量将是对合并存储的视图，并转换为目标设备。

示例

>>> data = TensorDict({"a": 1.0}, [], device=None) >>> data_cuda = data.to("cuda:0") # casts to cuda >>> data_int = data.to(torch.int) # casts to int >>> data_cuda_int = data.to("cuda:0", torch.int) # multiple casting >>> data_cuda = data.to(torch.randn(3, device="cuda:0")) # using an example tensor >>> data_cuda = data.to(other=TensorDict({}, [], device="cuda:0")) # using a tensordict example

to_dict() → dict[str, Any]¶

返回一个字典，其中键值对与 TensorDict 的键值对匹配。

to_h5(filename, **kwargs)¶

将 TensorDict 转换为使用 h5 后端的 PersistentTensorDict。

参数:

filename (str or path) – h5 文件的路径。

device (torch.device or compatible, optional) – 返回张量后，预期所在的设备。默认为 None（默认在 cpu 上）。

**kwargs – 要传递给 h5py.File.create_dataset() 的关键字参数。

返回值：:

与新创建的文件链接的 PersitentTensorDict 实例。

示例

>>> import tempfile >>> import timeit >>> >>> from tensordict import TensorDict, MemoryMappedTensor >>> td = TensorDict({ ... "a": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000)), ... "b": {"c": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000, 3))}, ... }, [1_000_000]) >>> >>> file = tempfile.NamedTemporaryFile() >>> td_h5 = td.to_h5(file.name, compression="gzip", compression_opts=9) >>> print(td_h5) PersistentTensorDict( fields={ a: Tensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float32, is_shared=False), b: PersistentTensorDict( fields={ c: Tensor(shape=torch.Size([1000000, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([1000000]), device=None, is_shared=False)}, batch_size=torch.Size([1000000]), device=None, is_shared=False)

to_module(module: nn.Module, *, inplace: bool | None = None, return_swap: bool = True, swap_dest=None, use_state_dict: bool = False, non_blocking: bool = False, memo=None)¶

将 TensorDictBase 实例的内容递归地写入给定 nn.Module 的属性。

参数:

module (nn.Module) – 要写入参数的模块。

关键字参数:

inplace (bool, optional) – 如果为 True，则模块中的参数或张量将被就地更新。默认值为 False。

return_swap (bool, optional) – 如果为 True，则将返回旧的参数配置。默认值为 False。

swap_dest (TensorDictBase, optional) – 如果 return_swap 为 True，则为应写入交换的 tensordict。

use_state_dict (bool, optional) – 如果为 True，则将使用 state-dict API 来加载参数（包括 state-dict 钩子）。默认值为 False。

non_blocking (bool, 可选) – 如果为 True 并且此副本位于不同的设备之间，则副本可能与主机异步发生。

示例

>>> from torch import nn >>> module = nn.TransformerDecoder( ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), ... num_layers=1) >>> params = TensorDict.from_module(module) >>> params.zero_() >>> params.to_module(module) >>> assert (module.layers[0].linear1.weight == 0).all()

to_namedtuple(dest_cls: type | None = None)¶

将 tensordict 转换为命名元组。

参数:

dest_cls (Type, optional) – 要使用的可选命名元组类。

示例

>>> from tensordict import TensorDict >>> import torch >>> data = TensorDict({ ... "a_tensor": torch.zeros((3)), ... "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3]) >>> data.to_namedtuple() GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))

to_padded_tensor(padding=0.0, mask_key: NestedKey | None = None)¶

将所有嵌套张量转换为填充版本，并相应地调整批次大小。

参数:

padding (float) – tensordict 中张量的填充值。默认值为 0.0。

mask_key (NestedKey, optional) – 如果提供，则为将写入有效值的掩码的键。如果异构维度不是 tensordict 批次大小的一部分，则会导致错误。默认值为 None

to_pytree()¶

将 tensordict 转换为 PyTree。

如果 tensordict 不是从 pytree 创建的，则此方法将返回 self 而不进行修改。

有关更多信息和示例，请参阅 from_pytree()。

to_tensordict() → T¶

从 TensorDictBase 返回一个常规 TensorDict 实例。

返回值：:

包含相同值的新的 TensorDict 对象。

transpose(dim0, dim1)¶

返回输入的转置版本 tensordict。给定维度 dim0 和 dim1 将被交换。

转置 tensordict 的就地或非就地修改也会影响原始 tensordict，因为内存是共享的，并且操作被映射回原始 tensordict。

示例

>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4]) >>> tensordict_transpose = tensordict.transpose(0, 1) >>> print(tensordict_transpose.shape) torch.Size([4, 3]) >>> tensordict_transpose.set("b",, torch.randn(4, 3)) >>> print(tensordict.get("b").shape) torch.Size([3, 4])

trunc() → T¶

计算 TensorDict 的每个元素的 trunc() 值。

trunc_() → T¶

就地计算 TensorDict 的每个元素的 trunc() 值。

type(dst_type)¶

将所有张量强制转换为 dst_type。

参数:

dst_type (type or string) – 目标类型

uint16()¶

将所有张量强制转换为 torch.uint16。

uint32()¶

将所有张量强制转换为 torch.uint32。

uint64()¶

将所有张量强制转换为 torch.uint64。

uint8()¶

将所有张量强制转换为 torch.uint8。

unbind(dim: int) → tuple[T, ...]¶

返回索引 tensordicts 的元组，沿指示的维度解绑。

示例

>>> td = TensorDict({ ... 'x': torch.arange(12).reshape(3, 4), ... }, batch_size=[3, 4]) >>> td0, td1, td2 = td.unbind(0) >>> td0['x'] tensor([0, 1, 2, 3]) >>> td1['x'] tensor([4, 5, 6, 7])

unflatten(dim, unflattened_size)¶

取消扁平化 tensordict 的维度，将其扩展到所需的形状。

参数:

dim (int) – 指定要取消扁平化的输入张量的维度。

unflattened_size (shape) – 是 tensordict 取消扁平化维度的新的形状。

示例

>>> td = TensorDict({ ... "a": torch.arange(60).view(3, 4, 5), ... "b": torch.arange(12).view(3, 4)}, ... batch_size=[3, 4]) >>> td_flat = td.flatten(0, 1) >>> td_unflat = td_flat.unflatten(0, [3, 4]) >>> assert (td == td_unflat).all()

unflatten_keys(separator: str = '.', inplace: bool = False) → T¶

将扁平的 tensordict 递归地转换为嵌套的 tensordict。

TensorDict 类型将丢失，结果将是一个简单的 TensorDict 实例。嵌套 tensordict 的元数据将从根推断：数据树中的所有实例将共享相同的批次大小、维度名称和设备。

参数:

separator (str, optional) – 嵌套项之间的分隔符。

inplace (bool, optional) – 如果为 True，则生成的 tensordict 将与调用发生所在的 tensordict 具有相同的标识。默认为 False。

示例

>>> data = TensorDict({"a": 1, "b - c": 2, "e - f - g": 3}, batch_size=[]) >>> data.unflatten_keys(separator=" - ") TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), e: TensorDict( fields={ f: TensorDict( fields={ g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)

此方法和 unflatten_keys() 在处理状态字典时特别有用，因为它们使您可以无缝地将扁平字典转换为模拟模型结构的数据结构。

示例

>>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4)) >>> ddp_model = torch.ao.quantization.QuantWrapper(model) >>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".") >>> print(state_dict) TensorDict( fields={ module: TensorDict( fields={ 0: TensorDict( fields={ bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> model_state_dict = state_dict.get("module") >>> print(model_state_dict) TensorDict( fields={ 0: TensorDict( fields={ bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> model.load_state_dict(dict(model_state_dict.flatten_keys(".")))

unlock_() → T¶

为非就地操作解锁 tensordict。

可以用作装饰器。

有关更多详细信息，请参阅 lock_()。

unsqueeze(dim: int) → T¶

为位于 -td.batch_dims 和 td.batch_dims 之间的维度解压缩所有张量，并将它们返回到一个新的 tensordict 中。

参数:

dim (int) – 解压缩的维度

示例

>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 4, 2), ... }, batch_size=[3, 4]) >>> td = td.unsqueeze(-2) >>> td.shape torch.Size([3, 1, 4]) >>> td.get("x").shape torch.Size([3, 1, 4, 2])

此操作也可以用作上下文管理器。对原始 tensordict 的更改将发生在原地，即原始张量的内容不会改变。这也假定 tensordict 未锁定（否则，需要解锁 tensordict）。

>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 4, 2), ... }, batch_size=[3, 4]) >>> with td.unsqueeze(-2) as tds: ... tds.set("y", torch.zeros(3, 1, 4)) >>> assert td.get("y").shape == [3, 4]

update(input_dict_or_td: dict[str, CompatibleType] | T, clone: bool = False, inplace: bool = False, *, non_blocking: bool = False, keys_to_update: Sequence[NestedKey] | None = None, is_leaf: Callable[[Type], bool] | None = None) → T¶

使用字典或另一个 TensorDict 中的值更新 TensorDict。

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。

clone (bool, optional) – 输入（张量）字典中的张量在设置之前是否应该克隆。默认为 False。

inplace (bool, optional) – 如果为 True，并且如果键与 tensordict 中的现有键匹配，则更新将针对该键值对在原地发生。如果找不到该条目，则会添加它。默认为 False。

关键字参数:

keys_to_update (NestedKeys 序列, optional) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update(data_src.select(*keys_to_update))。

non_blocking (bool, 可选) – 如果为 True 并且此副本位于不同的设备之间，则副本可能与主机异步发生。

is_leaf (Callable[[Type], bool], optional) – 一个可调用对象，指示对象类型是否被视为叶节点并被交换或张量集合。

返回值：:

self

示例

>>> td = TensorDict({}, batch_size=[3]) >>> a = torch.randn(3) >>> b = torch.randn(3, 4) >>> other_td = TensorDict({"a": a, "b": b}, batch_size=[]) >>> td.update(other_td, inplace=True) # writes "a" and "b" even though they can't be found >>> assert td['a'] is other_td['a'] >>> other_td = other_td.clone().zero_() >>> td.update(other_td) >>> assert td['a'] is not other_td['a']

update_(input_dict_or_td: dict[str, CompatibleType] | T, clone: bool = False, *, non_blocking: bool = False, keys_to_update: Sequence[NestedKey] | None = None) → T¶

使用字典或另一个 TensorDict 中的值在原地更新 TensorDict。

与 update() 不同，如果键对 self 未知，此函数将引发错误。

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。

clone (bool, optional) – 输入（张量）字典中的张量在设置之前是否应该克隆。默认为 False。

关键字参数:

keys_to_update (NestedKeys 序列, optional) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update_(data_src.select(*keys_to_update))。

non_blocking (bool, 可选) – 如果为 True 并且此副本位于不同的设备之间，则副本可能与主机异步发生。

返回值：:

self

示例

>>> a = torch.randn(3) >>> b = torch.randn(3, 4) >>> td = TensorDict({"a": a, "b": b}, batch_size=[3]) >>> other_td = TensorDict({"a": a*0, "b": b*0}, batch_size=[]) >>> td.update_(other_td) >>> assert td['a'] is not other_td['a'] >>> assert (td['a'] == other_td['a']).all() >>> assert (td['a'] == 0).all()

update_at_(input_dict_or_td: dict[str, CompatibleType] | T, idx: IndexType, clone: bool = False, *, non_blocking: bool = False, keys_to_update: Sequence[NestedKey] | None = None) → T¶

在指定索引处使用字典或另一个 TensorDict 中的值在原地更新 TensorDict。

与 TensorDict.update 不同，如果键对 TensorDict 未知，此函数将引发错误。

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。

idx (int, torch.Tensor, iterable, slice) – 应发生更新的 tensordict 的索引。

clone (bool, optional) – 输入（张量）字典中的张量在设置之前是否应该克隆。默认值为 False。

关键字参数:

keys_to_update (NestedKeys 序列, optional) – 如果提供，则仅更新 key_to_update 中的键列表。

non_blocking (bool, 可选) – 如果为 True 并且此副本位于不同的设备之间，则副本可能与主机异步发生。

返回值：:

self

示例

>>> td = TensorDict({ ... 'a': torch.zeros(3, 4, 5), ... 'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4]) >>> td.update_at_( ... TensorDict({ ... 'a': torch.ones(1, 4, 5), ... 'b': torch.ones(1, 4, 10)}, batch_size=[1, 4]), ... slice(1, 2)) TensorDict( fields={ a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32), b: Tensor(torch.Size([3, 4, 10]), dtype=torch.float32)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False) >>> assert (td[1] == 1).all()

values(include_nested: bool = False, leaves_only: bool = False, is_leaf=None) → Iterator[Tensor]¶

返回表示张量字典值的生成器。

参数:

include_nested (bool, optional) – 如果为 True，将返回嵌套值。默认为 False。

leaves_only (bool, optional) – 如果为 False，将只返回叶子。默认为 False。

is_leaf – 一个可选的 callable，用于指示某个类是否应被视为叶子。

var(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: bool | None = None) → TensorDictBase | torch.Tensor¶

返回输入张量字典中所有元素的方差值。

参数:

dim (int, tuple of int, optional) – 如果为 None，则返回一个无维张量字典，其中包含所有叶子的总和值（如果可以计算）。如果为整数或整数元组，则仅当该维度与张量字典形状兼容时，才会在指定的维度上调用 var。

keepdim (bool) – 输出张量是否保留维度。

关键字参数:

correction (int) – 样本大小与样本自由度之间的差值。默认为贝塞尔校正，correction=1。

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上进行缩减，并返回一个单个缩减后的张量。默认为 False。

view(*shape: int)¶

view(dtype)

view(shape: Size)

返回一个张量字典，其中包含根据新形状的张量视图，该形状与张量字典的 batch_size 兼容。

或者，可以提供 dtype 作为第一个未命名的参数。在这种情况下，所有张量都将以相应的 dtype 进行查看。请注意，这假设新形状与提供的 dtype 兼容。有关 dtype 视图的更多信息，请参见 view()。

参数:

*shape (int) – resulting tensordict 的新形状。

dtype (torch.dtype) – 可选地，一个用于表示张量内容的 dtype。

size – 可迭代对象

关键字参数:

batch_size (torch.Size, optional) – 如果提供了 dtype，则可以使用此关键字参数重置批次大小。如果使用形状调用 view，则此方法无效。

返回值：:

一个具有所需 batch_size 的新张量字典。

示例

>>> td = TensorDict(source={'a': torch.zeros(3,4,5), ... 'b': torch.zeros(3,4,10,1)}, batch_size=torch.Size([3, 4])) >>> td_view = td.view(12) >>> print(td_view.get("a").shape) # torch.Size([12, 5]) >>> print(td_view.get("b").shape) # torch.Size([12, 10, 1]) >>> td_view = td.view(-1, 4, 3) >>> print(td_view.get("a").shape) # torch.Size([1, 4, 3, 5]) >>> print(td_view.get("b").shape) # torch.Size([1, 4, 3, 10, 1])

where(condition, other, *, out=None, pad=None)¶

返回一个 TensorDict，其中包含从 self 或 other 中选择，具体取决于条件的元素。

参数:

condition (BoolTensor) – 当为 True（非零）时，生成 self，否则生成 other。

other (TensorDictBase or Scalar) – 值（如果 other 是标量）或在条件为 False 的索引处选择的元素。

关键字参数:

out (TensorDictBase, optional) – 输出 TensorDictBase 实例。

pad (scalar, optional) – 如果提供，则源或目标张量字典中缺少的键将被写入为 torch.where(mask, self, pad) 或 torch.where(mask, pad, other)。默认值为 None，即不允许缺少键。

zero_() → T¶

将张量字典中的所有张量归零（就地操作）。

zero_grad(set_to_none: bool = True) → T¶

递归地将 TensorDict 中所有张量的梯度归零。

参数:

set_to_none (bool, optional) – 如果为 True，则 tensor.grad 将为 None，否则为 0。默认值为 True。

TensorDictBase¶

文档

教程

资源