TensorDictParams¶

class tensordict.TensorDictParams(parameters: Optional[Union[TensorDictBase, dict]] = None, *, no_convert=False, lock: bool = False, **kwargs)¶

一个用于暴露参数的 TensorDict Base 类包装器。

此类设计用于持有包含参数的 TensorDictBase 实例，使其可被父级 Module 访问。这允许 tensordict 参数与 PyTorch 模块无缝集成，支持参数迭代和优化等操作。

主要特性

参数暴露：tensordict 内的参数会被暴露给父级模块，使其能够包含在 named_parameters() 等操作中。
索引：索引的操作方式与被包装的 tensordict 类似。但是，参数名称（在 named_parameters() 中）会使用 TensorDict.flatten_keys(“_”) 进行注册，这可能导致与 tensordict 内容不同的键名。
自动转换：在 tensordict 中设置的任何张量都会自动转换为 torch.nn.Parameter，除非通过 no_convert 关键字参数另行指定。

参数 (Args)

parameters (TensorDictBase or dict): 要表示为参数的 tensordict。值将被转换为: parameters，除非设置了 no_convert=True。如果提供了 dict，它将被包装在 TensorDict 实例中。也可以使用关键字参数。

关键字参数:

no_convert (bool) – 如果为 True，则不进行到 nn.Parameter 的转换，所有非参数、非缓冲张量将被转换为 Buffer 实例。如果为 False，则所有具有非整数 dtype 的张量将被转换为 Parameter，而整数 dtype 将被转换为 Buffer 实例。默认为 False。
lock (bool) –
如果为 True，则 TensorDictParams 托管的 tensordict 将被锁定，阻止修改，并在需要 unlock_() 时可能影响性能。默认为 False。

警告

由于内部 tensordict 默认不会被复制或锁定，因此在 TensorDictParams 中注册 tensordict 并随后修改其内容 __不会__ 更新 TensorDictParams parameters() 和 buffers() 序列中的值。
**kwargs – 用于填充 TensorDictParams 的键值对。与 parameters 输入互斥。

示例

>>> from torch import nn
>>> from tensordict import TensorDict
>>> module = nn.Sequential(nn.Linear(3, 4), nn.Linear(4, 4))
>>> params = TensorDict.from_module(module)
>>> params.lock_()
>>> p = TensorDictParams(params)
>>> print(p)
TensorDictParams(params=TensorDict(
    fields={
        0: TensorDict(
            fields={
                bias: Parameter(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                weight: Parameter(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        1: TensorDict(
            fields={
                bias: Parameter(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                weight: Parameter(shape=torch.Size([4, 4]), 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))
>>> class CustomModule(nn.Module):
...     def __init__(self, params):
...         super().__init__()
...         self.params = params
>>> m = CustomModule(p)
>>> # The wrapper supports assignment, and values are converted to Parameters
>>> m.params['other'] = torch.randn(3)
>>> assert isinstance(m.params['other'], nn.Parameter)

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

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

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

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

add(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, alpha: float | None = None, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

将 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, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

add_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor | float, *, alpha: float | None = None) → Any¶: 原地版本的 add()。

注意

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

add_module(name: str, module: Optional[Module]) → None¶

将子模块添加到当前模块。

可以使用给定的名称作为属性访问该模块。

参数:

name (str) – 子模块的名称。可以使用给定的名称从此模块访问子模块
module (Module) – 要添加到模块中的子模块。

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

执行 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 或 Tensor) – 被除数 tensordict（或 tensor）
tensor2 (TensorDict 或 Tensor) – 除数 tensordict（或 tensor）

关键字参数:

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

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

addcmul(other1: tensordict.base.TensorDictBase | torch.Tensor, other2: tensordict.base.TensorDictBase | torch.Tensor, *, value: float | None = 1) → Any¶

执行 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) – 要相乘的 tensordict 或 tensor
other2 (TensorDict 或 Tensor) – 要相乘的 tensordict 或 tensor

关键字参数:

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

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

all(dim: int | None = None) → bool | tensordict.base.TensorDictBase¶

检查 tensordict 中的所有值是否都为 True/非空。

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

amax(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None) → Self | torch.Tensor¶

返回输入 tensordict 中所有元素的最大值。

与 max() 相同，但 return_indices=False。

amin(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None) → Self | torch.Tensor¶

返回输入 tensordict 中所有元素的最小值。

与 min() 相同，但 return_indices=False。

any(dim: int | None = None) → bool | tensordict.base.TensorDictBase¶

检查 tensordict 中的任何值是否为 True/非空。

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

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

将一个可调用对象应用于tensordict中存储的所有值，并将它们设置在一个新的tensordict中。

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

参数:

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

关键字参数:

batch_size (sequence of int, optional) – 如果提供，则结果 TensorDict 将具有所需的 batch_size。 batch_size 参数应与转换后的 batch_size 匹配。这是一个仅限关键字的参数。
device (torch.device, 可选) – 生成的设备，如果存在。
names (字符串列表, 可选) – 新的维度名称，以防batch_size被修改。
inplace (bool, optional) – 如果为 True，则原地修改。默认为 False。这是一个仅限关键字的参数。
default (Any, 可选) – 其他tensordict中缺失条目的默认值。如果未提供，缺失的条目将引发KeyError。
filter_empty (bool, optional) – 如果为 True，则会过滤掉空的 tensordict。这也会降低计算成本，因为不会创建和销毁空的数据结构。非张量数据被视为叶节点，因此即使未被函数触及，也会保留在 tensordict 中。默认为 False 以向后兼容。
propagate_lock (bool, optional) – 如果为 True，则锁定的 tensordict 将产生另一个锁定的 tensordict。默认为 False。

call_on_nested (bool, optional) –

如果为 True，则该函数将应用于第一级张量和容器（TensorDict 或 tensorclass）。在这种情况下，func 负责传播其调用到嵌套级别。这允许在将调用传播到嵌套 tensordicts 时进行精细控制。如果为 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 执行的操作需要访问多个键来进行单次计算，则将 out 参数设置为 self 可能会导致操作产生静默错误的结果。例如

>>> 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) → Any¶

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

参数:

fn (Callable) – 要应用于tensordict中张量的函数。
*others (TensorDictBase 序列, 可选) – 要使用的其他tensordicts。

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

返回:: self 或 self 的副本，其中应用了函数

as_tensor() → Any¶: 将 tensordict 的每个叶子节点转换为普通的 torch.Tensor。

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

asin_() → Any¶: 原地计算TensorDict中每个元素的asin()值。

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

atan_() → Any¶: 原地计算TensorDict中每个元素的atan()值。

auto_batch_size_(batch_dims: Optional[int] = None, keep_compliant_size: bool = False) → Any¶

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

参数:

batch_dims (int, 可选) – 如果提供，则批次大小最多为 batch_dims。
keep_compliant_size (bool, 可选) – 如果为 True，则大小符合要求的子 tensordict 在 batch_dims 被传递时，其形状不会改变。如果为 False，所有包含的 tensordicts 的 batch_dims 都将与 batch_dims 匹配。默认为 False。

返回:

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])

auto_device_() → Any¶

自动设置设备，如果它是唯一的。

返回: self 及其编辑后的device属性。

property batch_dims: int¶

tensordict批次大小的长度。

返回:: 描述tensordict维度的整数。

property batch_size: Size¶

TensorDict的形状（或batch_size）。

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

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

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

返回:: 一个 Size 对象，描述 TensorDict 的批次大小。

示例

>>> 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() → Any¶: 将所有张量转换为torch.bfloat16。

bitwise_and(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

执行`self`和`other`之间的按位AND运算。

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

参数:: other (TensorDictBase 或 torch.Tensor) – 要执行按位 AND 的张量或 TensorDict。
关键字参数:: default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

bool() → Any¶: 将所有张量转换为torch.bool。

buffers(recurse: bool = True) → Iterator[Tensor]¶

返回模块缓冲区的迭代器。

参数:: recurse (bool) – 如果为 True，则会生成此模块以及所有子模块的缓冲区。否则，仅生成此模块的直接成员缓冲区。
产生:: torch.Tensor – 模块缓冲区

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

bytes(*, count_duplicates: bool = True) → int¶

计算包含的张量的字节数。

关键字参数:: count_duplicates (bool) – 是否计算重复张量为独立张量。如果为 False，则仅计算严格相同的张量（来自同一基张量的不同视图但具有不同 id 的张量将被计算两次）。默认为 True（每个张量被假定为单个副本）。

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

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

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

cat_from_tensordict(dim: int = 0, *, sorted: Optional[Union[bool, List[NestedKey]]] = None, out: Optional[Tensor] = None) → Tensor¶

将tensordict的所有条目连接成一个单一的张量。

参数:

dim (int, 可选) – 应沿哪个维度连接条目。

关键字参数:

sorted (bool 或 list of NestedKeys) – 如果为 True，则条目将按字母顺序连接。如果为 False（默认），则使用字典顺序。或者，可以提供一个键名列表，然后将张量按此顺序连接。这会产生一些开销，因为键列表将与 tensordict 中的叶节点名称列表进行比较。
out (torch.Tensor, 可选) – 连接操作的可选目标张量。

cat_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → Any¶

将条目连接成一个新条目，并可能删除原始值。

参数:

keys (NestedKey 序列) – 要连接的条目。

关键字参数:

out_key (NestedKey) – 连接输入的新键名。
keep_entries (bool, optional) – 如果为 False，则 keys 中的条目将被删除。默认为 False。
dim (int, optional) – 必须进行连接的维度。默认为 0。

返回: self

示例

>>> td = TensorDict(a=torch.zeros(1), b=torch.ones(1))
>>> td.cat_tensors("a", "b", out_key="c")
>>> assert "a" not in td
>>> assert (td["c"] == torch.tensor([0, 1])).all()

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

ceil_() → Any¶: 原地计算TensorDict中每个元素的ceil()值。

children() → Iterator[Module]¶

返回直接子模块的迭代器。

产生:: Module – 子模块

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

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

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

参数:

chunks (int) – 要返回的块数
dim (int, 可选) – 分割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(min: tensordict.base.TensorDictBase | torch.Tensor = None, max: tensordict.base.TensorDictBase | torch.Tensor = None, *, out=None) → Any¶

将 self 中的所有元素限制在 [ min, max ] 范围内。

令 min_value 和 max_value 分别为 min 和 max，则返回

\[y_i = \min(\max(x_i, \text{min\_value}_i), \text{max\_value}_i)\]

如果 min 为 None，则没有下限。或者，如果 max 为 None，则没有上限。

注意

如果 min 大于 max，则 torch.clamp(..., min, max) 将 input 中的所有元素设置为 max 的值。

clamp_max(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

如果`self`中的元素大于`other`，则将它们限制为`other`。

参数:: other (TensorDict 或 Tensor) – 另一个输入tensordict或张量。
关键字参数:: default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

clamp_max_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: 原地版本的 clamp_max()。

注意

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

clamp_min(other: tensordict.base.TensorDictBase | torch.Tensor, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

将 self 的元素限制在不小于 other 的值。

参数:: other (TensorDict 或 Tensor) – 另一个输入tensordict或张量。
关键字参数:: default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

clamp_min_(other: tensordict.base.TensorDictBase | torch.Tensor) → Any¶: 原地版本的 clamp_min()。

注意

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

clear() → Any¶: 清除 tensordict 的内容。

clear_device_() → Any¶

清除 tensordict 的设备。

返回: self

clear_refs_for_compile_() → Any¶

清除弱引用，以便 tensordict 安全地退出编译区域。

在返回 TensorDict 之前遇到 torch._dynamo.exc.Unsupported: reconstruct: WeakRefVariable() 时，请使用此方法。

返回: self

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

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

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

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

注意

与许多其他操作（逐元素算术、形状操作等）不同，clone 不会继承原始的锁定属性。此设计选择是为了能够创建一个可修改的副本，这是最常见的用法。

compile(*args, **kwargs)¶

使用 torch.compile() 编译此 Module 的前向传播。

此 Module 的 __call__ 方法将被编译，并且所有参数将按原样传递给 torch.compile()。

有关此函数的参数的详细信息，请参阅 torch.compile()。

complex128() → Any¶: 将所有张量转换为 torch.complex128。

complex32() → Any¶: 将所有张量转换为 torch.complex32。

complex64() → Any¶: 将所有张量转换为 torch.complex64。

consolidate(filename: Optional[Union[Path, str]] = None, *, num_threads=0, device: Optional[device] = 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 且提供了文件名，则会在共享内存中创建一个本地临时缓冲区，并将数据作为最后一步复制到存储位置。这可能比直接写入远程物理内存（例如 NFS）更快。默认为 False。
share_memory (bool, optional) – 如果为 True，则存储将放置在共享内存中。默认为 False。
pin_memory (bool, optional) – 是否将合并后的数据放入 pinned 内存。默认为 False。
metadata (bool, optional) – 如果为 True，则元数据将与通用存储一起存储。如果提供了文件名，则此参数无效。存储元数据在想要控制序列化方式时很有用，因为 TensorDict 在有或没有元数据可用时，对合并的 TD 的 pickling/unpickling 处理方式不同。

注意

如果 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())

contiguous(*args, **kwargs)¶: 返回一个相同类型的新 tensordict，其值是连续的（如果值已经是连续的，则返回 self）。(Returns a new tensordict of the same type with contiguous values (or self if values are already contiguous).)

copy() → Any¶

返回 tensordict 的浅拷贝（即，复制结构但不复制数据）。

等同于 TensorDictBase.clone(recurse=False)

copy_(tensordict: T, non_blocking: bool | None = None) → 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) → Any¶: 请参阅 TensorDictBase.update_at_。

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

cos_() → Any¶: 原地计算 TensorDict 中每个元素的 cos() 值。

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

cosh_() → Any¶: 原地计算 TensorDict 中每个元素的 cosh() 值。

cpu()¶

将 tensordict 转换为 CPU。

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

create_nested(key)¶

创建与当前 tensordict 具有相同形状、设备和维度名称的嵌套 tensordict。

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

示例

>>> 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=None)¶

将 tensordict 转换为 cuda 设备（如果尚未转换）。

参数:: device (int, optional) – 如果提供，则为张量应被转换到的 cuda 设备。

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

cummax(dim: int, *, reduce: Optional[bool] = None, return_indices: bool = True) → Union[Any, Tensor]¶

返回输入 tensordict 中所有元素的最大累积值。

参数:

dim (int) – 沿其执行 cummax 操作的维度的整数。

关键字参数:

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。
return_argmins (bool, optional) – 当传递 dim 参数时，cummax() 返回一个包含值和索引的命名元组。此功能的 TensorDict 等价功能是返回一个具有相同结构的张量类，其中包含条目 "values" 和 "indices"。默认为 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummax(dim=0)
cummax(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummax(reduce=True, dim=0)
torch.return_types.cummax(...)

cummin(dim: int, *, reduce: Optional[bool] = None, return_indices: bool = True) → Union[Any, Tensor]¶

返回输入 tensordict 中所有元素的最小累积值。

参数:

dim (int) – 沿其执行 cummin 操作的维度的整数。

关键字参数:

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。
return_argmins (bool, optional) – 当传递 dim 参数时，cummin() 返回一个包含值和索引的命名元组。此功能的 TensorDict 等价功能是返回一个具有相同结构的张量类，其中包含条目 "values" 和 "indices"。默认为 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummin(dim=0)
cummin(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummin(reduce=True, dim=0)
torch.return_types.cummin(...)

property data¶: 返回一个包含叶张量的 .data 属性的 tensordict。

data_ptr(*, storage: bool = False)¶

返回tensordict叶子节点的data_ptr。

这有助于检查两个tensordict是否共享相同的data_ptr()。

关键字参数:: storage (bool, optional) – 如果为 True，则调用 tensor.untyped_storage().data_ptr()。默认为 False。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> assert (td0.data_ptr() == td.data_ptr()).all()

注意

LazyStackedTensorDict 实例将显示为嵌套的 tensordicts，以反映其叶节点的实际 data_ptr()。

>>> td0 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> td1 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> td = TensorDict.lazy_stack([td0, td1])
>>> td.data_ptr()
TensorDict(
    fields={
        0: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False),
        1: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

del_(*args, **kwargs)¶

删除tensordict的键。

参数:: key (NestedKey) – 要删除的键
返回:: self

densify(layout: layout = torch.strided)¶

尝试用连续张量（普通张量或嵌套张量）来表示懒惰堆栈。

关键字参数:: layout (torch.layout) – 嵌套张量的布局，如果有的话。默认为 strided。

property depth: int¶

返回tensordict的深度 - 最大层级数。

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

detach() → Any¶

分离tensordict中的张量。

返回:: 返回一个不包含需要梯度的张量的新tensordict。

detach_(*args, **kwargs)¶

就地分离tensordict中的张量。

返回:: self。

property device¶

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: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

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

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

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

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

div_(other: tensordict.base.TensorDictBase | torch.Tensor) → Any¶: 原地执行的 div() 版本。

注意

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

double() → Any¶: 将所有张量转换为torch.bool。

property dtype¶: 返回tensordict中值的dtype，如果它是唯一的。

dumps(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → Any¶

将tensordict保存到磁盘。

此函数是 memmap() 的代理。

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

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

参数:

recurse (bool, optional) – 如果为 True，则将复制 TensorDict 的整个结构而没有内容。否则，仅复制根。默认为 False。

关键字参数:

batch_size (torch.Size, optional) – tensordict 的新批次大小。
device (torch.device, optional) – 新设备。
names (list of str, optional) – 维度名称。

entry_class(*args, **kwargs)¶

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

当 get() 执行可能昂贵时，此方法应优于 tensordict.get(key).shape。

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

erf_() → Any¶: 就地计算TensorDict中每个元素的erf()值。

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

erfc_() → Any¶: 就地计算TensorDict中每个元素的erfc()值。

eval() → Self¶

将模块设置为评估模式。

这仅对某些模块有影响。有关模块在训练/评估模式下的行为，例如它们是否受影响（如 Dropout、BatchNorm 等），请参阅具体模块的文档。

这等同于 self.train(False)。

有关 .eval() 和几种可能与之混淆的类似机制之间的比较，请参阅局部禁用梯度计算。

返回:: self
返回类型:: 模块

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

排除tensordict的键，并返回一个不包含这些条目的新tensordict。

值不会被复制：对原始tensordict或新tensordict的张量的就地修改将导致两个tensordict都发生变化。

参数:

*keys (str) – 要排除的键。
inplace (bool) – 如果为 True，则 tensordict 将就地剪枝。默认为 False。

返回:

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

示例

>>> 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() → Any¶: 计算TensorDict中每个元素的exp()值。

exp_() → Any¶: 就地计算TensorDict中每个元素的exp()值。

expand(*args, **kwargs) → Any¶

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

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

示例

>>> 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: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶

将tensordict的形状广播到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() → Any¶: 计算TensorDict中每个元素的expm1()值。

expm1_() → Any¶: 就地计算TensorDict中每个元素的expm1()值。

extra_repr() → str¶

返回模块的额外表示。

要打印自定义额外信息，您应该在自己的模块中重新实现此方法。单行和多行字符串均可接受。

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

Fills a tensor pointed by the key with a given scalar value.

参数:

key (str or nested key) – entry to be filled.
value (Number or bool) – value to use for the filling.

返回:

self

filter_empty_()¶: 就地过滤掉所有空的tensordict。

filter_non_tensor_data() → Any¶: 过滤掉所有非张量数据。

flatten(start_dim: int | None = None, end_dim: int | None = None)¶

展平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) → TensorDictBase¶

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

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

参数:

separator (str, optional) – 嵌套项之间的分隔符。
inplace (bool, optional) – 如果为 True，则结果 tensordict 将与调用它的 tensordict 具有相同的身份。默认为 False。
is_leaf (callable, optional) –
一个作用于类类型的可调用对象，返回一个布尔值，指示该类是否应被视为叶子节点。

注意

is_leaf 的目的是不是阻止递归调用嵌套的 tensordicts，而是为过滤标记某些类型，当 `leaves_only=True` 时。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，嵌套 tensordict 的结构仍将被遍历。换句话说，`is_leaf` 不控制递归深度，而是提供了一种当 `leaves_only=True` 时从结果中过滤掉某些类型的方法。这意味着树中的一个节点既可以是叶节点，也可以是具有子节点的节点。实际上，`is_leaf` 的默认值不包含 tensordict 和 tensorclass 实例作为叶节点。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。

示例

>>> 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() 方法在处理 state-dicts 时尤其有用，因为它们可以无缝地将扁平字典转换为模仿模型结构的数数据结构。

示例

>>> 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() → Any¶: 将所有张量转换为torch.float。

float16() → Any¶: 将所有张量转换为torch.float16。

float32() → Any¶: 将所有张量转换为torch.float32。

float64() → Any¶: 将所有张量转换为torch.float64。

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

floor_() → Any¶: 就地计算TensorDict中每个元素的floor()值。

forward(*input: Any) → None¶

定义每次调用时执行的计算。

所有子类都应重写此方法。

注意

尽管前向传播的实现需要在此函数中定义，但您应该在之后调用 Module 实例而不是此函数，因为前者会处理注册的钩子，而后者则会静默忽略它们。

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

frac_() → Any¶: 就地计算TensorDict中每个元素的frac()值。

classmethod from_any(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

Recursively converts any object to a TensorDict.

注意

from_any 比常规的 TensorDict 构造函数限制更少。它可以使用自定义启发式方法将数据结构（如 dataclasses 或 tuples）转换为 tensordict。此方法可能会产生一些额外的开销，并在映射策略方面涉及更多主观选择。

注意

This method recursively converts the input object to a TensorDict. If the object is already a TensorDict (or any similar tensor collection object), it will be returned as is.

参数:

obj – The object to be converted.

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，将自动计算批次大小。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None（每个级别都有完整的批次大小）。
device (torch.device, optional) – 将创建 TensorDict 的设备。
batch_size (torch.Size, optional) – TensorDict 的批次大小。与 auto_batch_size 互斥。

返回:

A TensorDict representation of the input object.

Supported objects

通过 from_dataclass() 的数据类（数据类将被转换为 TensorDict 实例，而不是 tensorclasses）。
通过 from_namedtuple() 的命名元组。
通过 from_dict() 的字典。
通过 from_tuple() 的元组。
通过 from_struct_array() 的 NumPy 结构化数组。
通过 from_h5() 的 HDF5 对象。

classmethod from_dataclass(dataclass, *, dest_cls: Optional[Type] = None, auto_batch_size: bool = False, batch_dims: Optional[int] = None, as_tensorclass: bool = False, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

Converts a dataclass into a TensorDict instance.

参数:

dataclass – The dataclass instance to be converted.

关键字参数:

dest_cls (tensorclass, optional) – 用于映射数据的 tensorclass 类型。如果未提供，则创建一个新类。如果 `obj` 是一个类型或 as_tensorclass 为 `False`，则无效。
auto_batch_size (bool, optional) – 如果为 True，则自动确定并应用批次大小到生成的 TensorDict。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None（每个级别都有完整的批次大小）。
as_tensorclass (bool, optional) – 如果为 True，则将转换委托给函数 from_dataclass()，并返回一个兼容张量的类（tensorclass()）或实例，而不是 TensorDict。默认为 False。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

A TensorDict instance derived from the provided dataclass, unless as_tensorclass is True, in which case a tensor-compatible class or instance is returned.

抛出:

TypeError – If the provided input is not a dataclass instance.

警告

This method is distinct from the free function from_dataclass and serves a different purpose. While the free function returns a tensor-compatible class or instance, this method returns a TensorDict instance.

注意

此方法创建一个新的 TensorDict 实例，其键对应于输入 dataclass 的字段。
结果 TensorDict 中的每个键都使用 `cls.from_any` 方法进行初始化。
auto_batch_size 选项允许自动确定批次大小并将其应用于结果 TensorDict。

classmethod from_dict(*args, **kwargs)¶

从字典或另一个 TensorDict 创建 TensorDict。

If batch_size is not specified, returns the maximum batch size possible.

This function works on nested dictionaries too, or can be used to determine the batch-size of a nested tensordict.

参数:

input_dict (dictionary, optional) – a dictionary to use as a data source (nested keys compatible).

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，将自动计算批次大小。默认为 False。
batch_size (iterable of int, optional) – a batch size for the tensordict.
device (torch.device 或 兼容类型, optional) – TensorDict 的设备。
batch_dims (int, optional) – batch_dims（即用于 batch_size 的前导维度数量）。与 batch_size 互斥。注意，这是 tensordict 的__最大__批次维度数，较小的数字是允许的。
names (list of str, optional) – the dimension names of the 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, *, auto_batch_size: bool = False, batch_size=None, device=None, batch_dims=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)

classmethod from_h5(filename, *, mode: str = 'r', auto_batch_size: bool = False, batch_dims: Optional[int] = None, batch_size: Optional[Size] = None)¶

从 h5 文件创建 PersistentTensorDict。

参数:

filename (str) – h5 文件的路径。

关键字参数:

mode (str, 可选) – 读取模式。默认为 "r"。
auto_batch_size (bool, optional) – 如果为 True，将自动计算批次大小。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None（每个级别都有完整的批次大小）。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入 h5 文件的 PersistentTensorDict 表示。

示例

>>> td = TensorDict.from_h5("path/to/file.h5")
>>> print(td)
PersistentTensorDict(
    fields={
        key1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        key2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

classmethod from_module(module, as_module: bool = False, lock: bool = True, use_state_dict: bool = False)¶

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

参数:

module (nn.Module) – 要从中获取参数的模块。
as_module (bool, optional) – 如果为 True，则将返回一个 TensorDictParams 实例，该实例可用于将参数存储在 torch.nn.Module 中。默认为 False。
lock (bool, optional) – 如果为 True，则结果 tensordict 将被锁定。默认为 True。
use_state_dict (bool, optional) –
如果为 True，则将使用模块的 state-dict 并将其解压到具有模型树结构的 TensorDict 中。默认为 False。

注意

这在使用 state-dict hook 时尤其有用。

示例

>>> 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)

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

为 vmap 的 ensemable 学习/特征期望应用检索多个模块的参数。

参数:

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

关键字参数:

as_module (bool, optional) – 如果为 True，则返回一个 TensorDictParams 实例，可用于在 torch.nn.Module 中存储参数。默认为 False。
lock (bool, optional) – 如果为 True，则结果 tensordict 将被锁定。默认为 True。
use_state_dict (bool, optional) –
如果为 True，则将使用模块的 state-dict 并将其解压到具有模型树结构的 TensorDict 中。默认为 False。

注意

这在使用 state-dict hook 时尤其有用。
lazy_stack (bool, optional) –
是否密集堆叠或懒惰堆叠参数。默认为 False（密集堆叠）。

注意

lazy_stack 和 as_module 是互斥的特性。

警告

懒惰输出和非懒惰输出之间有一个重要的区别：非懒惰输出将使用所需的批次大小重新实例化参数，而 `lazy_stack` 将仅将参数表示为懒惰堆叠。这意味着，虽然原始参数可以安全地传递给优化器（当 `lazy_stack=True` 时），但在设置为 `True` 时需要传递新参数。

警告

虽然使用 lazy stack 来保留原始参数引用很有诱惑力，但请记住，lazy stack 在每次调用 get() 时都会执行堆叠操作。这将需要内存（参数大小的 N 倍，如果构建了图则更多）和计算时间。这也意味着优化器将包含更多参数，并且类似 step() 或 zero_grad() 的操作执行时间会更长。通常，lazy_stack 应仅用于少数用例。
expand_identical (bool, optional) – 如果为 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

classmethod from_namedtuple(named_tuple, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

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

参数:

named_tuple – 要转换的命名元组实例。

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，将自动计算批次大小。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None（每个级别都有完整的批次大小）。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入命名元组的 TensorDict 表示。

示例

>>> 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: Optional[Size] = None, auto_batch_size: bool = False, batch_dims: Optional[int] = None)¶

将 pytree 转换为 TensorDict 实例。

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

其他非张量键将被添加，以跟踪每个级别的标识，从而提供内置的 pytree 到 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 from_remote_init(src: int, group: 'ProcessGroup' | None = None, device: torch.device | None = None) → Self¶

从远程发送的元数据创建新的 tensordict 实例。

此类方法接收由 init_remote 发送的元数据，创建具有匹配形状和 dtype 的新 tensordict，然后异步接收实际的 tensordict 内容。

参数:

src (int) – 发送元数据的源进程的秩。
group ("ProcessGroup", optional) – 要使用的进程组。默认为 None。
device (torch.device, 可选) – 用于张量运算的设备。默认为 None。

返回:

使用接收到的元数据和内容初始化的新 tensordict 实例。

返回类型:

TensorDict

另请参阅

发送进程应已调用 ~.init_remote 来发送元数据和内容。

classmethod from_struct_array(struct_array: ndarray, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None) → Any¶

将结构化 numpy 数组转换为 TensorDict。

生成的 TensorDict 将与 numpy 数组共享相同的内存内容（这是一次零拷贝操作）。原地更改结构化 numpy 数组的值会影响 TensorDict 的内容。

注意

此方法执行零拷贝操作，这意味着生成的 TensorDict 将与输入的 numpy 数组共享相同的内存内容。因此，原地更改 numpy 数组的值会影响 TensorDict 的内容。

参数:

struct_array (np.ndarray) – 要转换的结构化 numpy 数组。

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，将自动计算批次大小。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None（每个级别都有完整的批次大小）。
device (torch.device, 可选) –
将创建 TensorDict 的设备。默认为 None。

注意

更改设备（即，指定任何非 `None` 或 `"cpu"` 的设备）将传输数据，从而导致返回数据的内存位置发生更改。
batch_size (torch.Size, 可选) – TensorDict 的批次大小。默认为 None。

返回:

输入的结构化 numpy 数组的 TensorDict 表示。

示例

>>> x = np.array(
...     [("Rex", 9, 81.0), ("Fido", 3, 27.0)],
...     dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")],
... )
>>> td = TensorDict.from_struct_array(x)
>>> x_recon = td.to_struct_array()
>>> assert (x_recon == x).all()
>>> assert x_recon.shape == x.shape
>>> # Try modifying x age field and check effect on td
>>> x["age"] += 1
>>> assert (td["age"] == np.array([10, 4])).all()

classmethod from_tuple(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

将元组转换为 TensorDict。

参数:

obj – 要转换的元组实例。

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，将自动计算批次大小。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None（每个级别都有完整的批次大小）。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入的元组的 TensorDict 表示。

示例

>>> my_tuple = (1, 2, 3)
>>> td = TensorDict.from_tuple(my_tuple)
>>> print(td)
TensorDict(
    fields={
        0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

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

从键列表和单个值创建 tensordict。

参数:

keys (list of NestedKey) – 指定新字典键的可迭代对象。
value (compatible type, optional) – 所有键的值。默认为 0。

gather(dim: int, index: Tensor, out: Optional[T] = None) → Any¶

沿由 dim 指定的轴收集值。

参数:

dim (int) – 要收集元素的维度
index (torch.Tensor) – 一个长整型张量，其维度数量与 tensordict 匹配，只有一个维度不同（收集维度）。它的元素引用沿所需维度收集的索引。
out (TensorDictBase, 可选) – 目标 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: 'torch.distributed.ProcessGroup' | None = None) → Self | None¶

从各个工作节点收集 tensordicts 并将它们堆叠到目标节点上的 self 中。

参数:

dst (int) – 目标工作节点的 rank，gather_and_stack() 将在该节点上调用。
group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `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://: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://: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 = None) → Union[Tensor, TensorCollection]¶

获取输入键对应的存储值。

参数:

key (str, str 的元组) – 要查询的键。如果是 str 的元组，则等同于链式调用 getattr。
default –
如果 tensordict 中找不到该键，则返回默认值。默认为 None。

警告

以前，如果 tensordict 中不存在某个键且未提供默认值，则会引发 `KeyError`。从 v0.7 开始，此行为已更改，并返回 `None` 值（符合 dict.get 的行为）。要采用旧行为，请设置环境变量 `export TD_GET_DEFAULTS_TO_NONE=’0’` 或调用 :func`~tensordict.set_get_defaults_to_none(False)`。

注意

当处理不规则张量时，可以将关键字参数传递给 get()。有关完整概述，请参阅 get()。

示例

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

get_at(key: NestedKey, *args, **kwargs) → Union[Tensor, TensorCollection]¶

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

参数:

key (str, tuple of str) – 要检索的键。
index (int, slice, torch.Tensor, 可迭代对象) – 张量的索引。
default (torch.Tensor) – 如果 tensordict 中不存在该键，则返回的默认值。

返回:

索引的张量。

示例

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

get_buffer(target: str) → Tensor¶

返回由 target 给定的缓冲区（如果存在），否则抛出错误。

有关此方法功能的更详细解释以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数:: target – 要查找的 buffer 的完全限定字符串名称。（要指定完全限定字符串，请参阅 get_submodule。）
返回:: 由 target 引用的缓冲区
返回类型:: torch.Tensor
抛出:: AttributeError – 如果目标字符串引用了无效路径或解析为非 buffer 的内容

get_extra_state() → Any¶

返回要包含在模块 state_dict 中的任何额外状态。

如果需要存储额外的状态，请实现此函数以及对应的 set_extra_state()。构建模块的 state_dict() 时会调用此函数。

注意，为了保证 state_dict 的序列化工作正常，额外状态应该是可被 pickle 的。我们仅为 Tensors 的序列化提供向后兼容性保证；其他对象的序列化形式若发生变化，可能导致向后兼容性中断。

返回:: 要存储在模块 state_dict 中的任何额外状态
返回类型:: 对象

get_item_shape(key: NestedKey)¶: 返回条目的形状，可能避免调用 get()。

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

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

此方法对张量/TensorDict 值具有鲁棒性，这意味着如果收集到的值是常规张量，它也会被返回（尽管此方法带有一些开销，不应超出其自然范围使用）。

有关如何在 tensordict 中设置非张量值的更多信息，请参阅 `tensordict.TensorDictBase.set_non_tensor`。

参数:

key (NestedKey) – 非张量数据的存储位置。
default (Any, optional) – 找不到键时要返回的值。

返回： `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)

get_parameter(target: str) → Parameter¶

如果存在，返回由 target 给定的参数，否则抛出错误。

有关此方法功能的更详细解释以及如何正确指定 target，请参阅 get_submodule 的文档字符串。

参数:: target – 要查找的 Parameter 的完全限定字符串名称。（要指定完全限定字符串，请参阅 get_submodule。）
返回:: 由 target 引用的参数
返回类型:: torch.nn.Parameter
抛出:: AttributeError – 如果目标字符串引用了无效路径或解析为非 nn.Parameter 的内容

get_submodule(target: str) → Module¶

如果存在，返回由 target 给定的子模块，否则抛出错误。

例如，假设您有一个 nn.Module A，它看起来像这样

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(图示了一个 nn.Module A。 A 包含一个嵌套子模块 net_b，该子模块本身有两个子模块 net_c 和 linear。 net_c 随后又有一个子模块 conv。)

要检查是否存在 linear 子模块，可以调用 get_submodule("net_b.linear")。要检查是否存在 conv 子模块，可以调用 get_submodule("net_b.net_c.conv")。

get_submodule 的运行时复杂度受 target 中模块嵌套深度的限制。与 named_modules 的查询相比，后者的复杂度是按传递模块数量计算的 O(N)。因此，对于简单地检查某个子模块是否存在，应始终使用 get_submodule。

参数:: target – 要查找的子模块的完全限定字符串名称。（要指定完全限定字符串，请参阅上面的示例。）
返回:: 由 target 引用的子模块
返回类型:: torch.nn.Module
抛出:: AttributeError – 如果沿目标字符串解析路径的任何一点解析为不存在的属性名或不是 nn.Module 实例的对象。

property grad¶: 返回一个 tensordict，其中包含叶子张量的 .grad 属性。

half() → Any¶: 将所有张量转换为 torch.half。

init_remote(dst: int, group: 'ProcessGroup' | None = None, device: torch.device | None = None) → None¶

通过发送元数据和内容来初始化远程 tensordict。

此方法将当前 tensordict 的元数据（形状、dtype 等）发送到指定的目标 rank（dst）。

然后异步发送实际的 tensordict 内容。

参数:

dst (int) – 目标进程的 rank。
group ("ProcessGroup", optional) – 要使用的进程组。默认为 None。
device (torch.device, 可选) – 用于张量运算的设备。默认为 None。

另请参阅

接收进程应调用 ~.from_remote_init 或等效方法来接收并基于发送的元数据初始化新的 tensordict。

示例

>>> import os
>>> import torch
>>> import torch.distributed as dist
>>> from tensordict import TensorDict, MemoryMappedTensor
>>> import multiprocessing as mp
>>>
>>> def server(queue):
...     # Set environment variables for distributed communication
...     os.environ["MASTER_ADDR"] = "localhost"
...     os.environ["MASTER_PORT"] = "29505"
...
...     # Initialize the distributed backend
...     dist.init_process_group("gloo", rank=0, world_size=2)
...
...     # Create a sample tensordict
...     td = (
...         TensorDict(
...             {
...                 ("a", "b"): torch.ones(2),
...                 "c": torch.ones(2),
...                 ("d", "e", "f"): MemoryMappedTensor.from_tensor(torch.ones(2, 2)),
...             },
...             [2],
...         )
...         .expand(1, 2)
...         .contiguous()
...     )
...
...     # Send the tensordict metadata and content to the client
...     td.init_remote(dst=1)
...
>>> def client(queue):
...     # Set environment variables for distributed communication
...     os.environ["MASTER_ADDR"] = "localhost"
...     os.environ["MASTER_PORT"] = "29505"
...
...     # Initialize the distributed backend
...     dist.init_process_group("gloo", rank=1, world_size=2)
...
...     # Receive the tensordict metadata and content from the server
...     received_td = TensorDict.from_remote_init(src=0)
...
...     # Verify that the received tensordict matches the expected structure and values
...     assert set(received_td.keys()) == {"a", "c", "d"}
...     assert (received_td == 1).all()
...
...     # Signal that the test has completed successfully
...     queue.put("yuppie")
>>>
>>> if __name__ == "__main__":
...     queue = mp.Queue(1)
...
...     # Create and start the server and client processes
...     main_worker = mp.Process(target=server, args=(queue,))
...     secondary_worker = mp.Process(target=client, args=(queue,))
...
...     main_worker.start()
...     secondary_worker.start()
...
...     try:
...         out = queue.get(timeout=10)  # Wait for the signal with a timeout
...         print(out)  # Should print "yuppie"
...     finally:
...         queue.close()
...         main_worker.join(timeout=10)
...         secondary_worker.join(timeout=10)

int() → Any¶: 将所有张量转换为 torch.int。

int16() → Any¶: 将所有张量转换为 torch.int16。

int32() → Any¶: 将所有张量转换为 torch.int32。

int64() → Any¶: 将所有张量转换为 torch.int64。

int8() → Any¶: 将所有张量转换为 torch.int8。

ipu(device: Optional[Union[device, int]] = None) → Self¶

将所有模型参数和缓冲区移动到 IPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 IPU 上，则应在构建优化器之前调用它。

注意

此方法就地修改模块。

参数:: device (int, optional) – 如果指定，所有参数都将复制到该设备
返回:: self
返回类型:: 模块

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

异步接收 tensordict 的内容并用其更新内容。

请参阅 isend() 方法中的示例以获取上下文。

参数:

src (int) – 源工作进程的 rank。

关键字参数:

group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `None`。
return_premature (bool) – 如果为 True，则返回一个 future 列表，在 tensordict 更新之前等待它们。默认为 False，即在调用期间等待更新完成。
init_tag (int) – 用于标记张量的 init_tag。注意，这将根据 TensorDict 中包含的张量数量进行递增。
pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许从不同节点发送多个数据而不会重叠。请注意，这些伪随机数的生成成本很高（每秒 1e-5），这意味着它可能会减慢算法的运行速度。此值必须与传递给 isend() 的值匹配。默认为 False。

返回:

如果 return_premature=True，则返回一个未来列表以等待直到 tensordict 更新。: upon until the tensordict is updated.

is_consolidated()¶: 检查 TensorDict 是否具有合并的存储。

is_contiguous(*args, **kwargs)¶: 返回一个布尔值，指示所有张量是否是连续的。

is_empty() → bool¶: 检查 tensordict 是否包含任何叶子。

is_floating_point() → bool¶: 检查 tensordict 中的所有张量是否为浮点数。

property is_memmap: bool¶

检查 tensordict 是否为内存映射。

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

对于 CUDA 设备上的 tensordict，这始终为 True。

property is_shared: bool¶

检查 tensordict 是否在共享内存中。

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

对于 CUDA 设备上的 tensordict，这始终为 True。

isend(dst: int, *, group: 'torch.distributed.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False, return_early: bool = False) → int | List['Work']¶

异步发送 tensordict 的内容。

参数:

dst (int) – 应将内容发送到的目标工作进程的 rank。

关键字参数:

group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `None`。
init_tag (int) – 用于标记张量的初始 init_tag。注意，此值将根据 TensorDict 中包含的张量数量进行递增。
pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许从不同节点发送多个数据而不会重叠。请注意，这些伪随机数的生成成本很高（每秒 1e-5），这意味着它可能会减慢算法的运行速度。默认为 False。
return_early (bool, optional) – 如果为 True，则返回一个 futures 列表，而不是最后一个发送的张量的 tag。默认为 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://: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://: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() → Any¶

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

实数值在非 NaN、负无穷或无穷大时是有限的。复数值在其实部和虚部都有限时是有限的。

isnan() → Any¶

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

当复数的实部和/或虚部为 NaN 时，复数值被视为 NaN。

isneginf() → Any¶: 测试输入中的每个元素是否为负无穷。

isposinf() → Any¶: 测试输入中的每个元素是否为负无穷。

isreal() → Any¶: 返回一个新的 tensordict，其中包含布尔值元素，表示输入中的每个元素是否为实值。

items(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) → Iterator[Union[Tensor, TensorCollection]]¶

返回 tensordict 的键值对生成器。

参数:

include_nested (bool, 可选) – 如果为 `True`，则返回嵌套值。默认为 `False`。
leaves_only (bool, 可选) – 如果为 `False`，则仅返回叶子节点。默认为 `False`。
is_leaf (callable, optional) –
一个作用于类类型的可调用对象，返回一个布尔值，指示该类是否应被视为叶子节点。

注意

is_leaf 的目的是不是阻止递归调用嵌套的 tensordicts，而是为过滤标记某些类型，当 `leaves_only=True` 时。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，嵌套 tensordict 的结构仍将被遍历。换句话说，`is_leaf` 不控制递归深度，而是提供了一种当 `leaves_only=True` 时从结果中过滤掉某些类型的方法。这意味着树中的一个节点既可以是叶节点，也可以是具有子节点的节点。实际上，`is_leaf` 的默认值不包含 tensordict 和 tensorclass 实例作为叶节点。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。

关键字参数:

sort (bool, optional) – 是否应该对键进行排序。对于嵌套键，键将根据其连接的名称进行排序（即，("a", "key") 将被计为 "a.key" 以进行排序）。请注意，当处理大型 tensordicts 时，排序可能会产生显著的开销。默认为 False。

keys(*args, **kwargs)¶

返回 tensordict 键的生成器。

警告

TensorDict 的 `keys()` 方法返回键的懒惰视图。如果查询了 `keys` 但未迭代，然后修改了 tensordict，则稍后迭代 `keys` 将返回新配置的键。

参数:

include_nested (bool, 可选) – 如果为 `True`，则返回嵌套值。默认为 `False`。
leaves_only (bool, 可选) – 如果为 `False`，则仅返回叶子节点。默认为 `False`。
is_leaf (callable, optional) –
一个作用于类类型的可调用对象，返回一个布尔值，指示该类是否应被视为叶子节点。

注意

is_leaf 的目的是不是阻止递归调用嵌套的 tensordicts，而是为过滤标记某些类型，当 `leaves_only=True` 时。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，嵌套 tensordict 的结构仍将被遍历。换句话说，`is_leaf` 不控制递归深度，而是提供了一种当 `leaves_only=True` 时从结果中过滤掉某些类型的方法。这意味着树中的一个节点既可以是叶节点，也可以是具有子节点的节点。实际上，`is_leaf` 的默认值不包含 tensordict 和 tensorclass 实例作为叶节点。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。

关键字参数:

sort (bool, optional) – 是否应该对键进行排序。对于嵌套键，键将根据其连接的名称进行排序（即，("a", "key") 将被计为 "a.key" 以进行排序）。请注意，当处理大型 tensordicts 时，排序可能会产生显著的开销。默认为 False。

示例

>>> 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: int = 0, *, out=None, **kwargs)¶

创建 TensorDicts 的懒惰堆叠。

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

lerp(end: tensordict._tensorcollection.TensorCollection | torch.Tensor, weight: tensordict._tensorcollection.TensorCollection | torch.Tensor | float) → Any¶

根据标量或张量 `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, 张量或float) – 插值公式的权重。

lerp_(end: tensordict.base.TensorDictBase | torch.Tensor | float, weight: tensordict.base.TensorDictBase | torch.Tensor | float)¶: 原地版本的 lerp()。

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

lgamma_() → Any¶: 原地计算 TensorDict 中每个元素的 lgamma() 值。

classmethod load(prefix: str | pathlib.Path, *args, **kwargs) → Any¶

从磁盘加载 tensordict。

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

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

在当前 tensordict 中从磁盘加载 tensordict。

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

classmethod load_memmap(prefix: str | pathlib.Path, device: Optional[device] = None, non_blocking: bool = False, *, out: Optional[TensorDictBase] = None) → Any¶

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

参数:

prefix (str 或 文件夹路径) – 应从中获取已保存 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()

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

示例

>>> 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 | pathlib.Path)¶

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

有关更多信息，请参阅 `tensordict.TensorDictBase.load_memmap`。

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

将格式与此 tensordict 的 state_dict() 相同的 state-dict 加载到 tensordict 中。

参数:

state_dict (OrderedDict) – 要复制的 state_dict。
strict (bool, optional) – 是否严格强制 state_dict 中的键与此 tensordict 的 torch.nn.Module.state_dict() 函数返回的键匹配。默认为 True。
assign (bool, optional) – 是否将 state 字典中的项分配给 tensordict 中相应的键，而不是就地复制到 tensordict 的当前张量中。当 False 时，将保留当前模块张量的属性；当 True 时，将保留 state 字典中张量的属性。默认为 False。
from_flatten (bool, optional) – 如果为 True，则假定输入的 state_dict 是扁平化的。默认为 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_() → Any¶

锁定 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() → Any¶: 计算 TensorDict 中每个元素的 log() 值。

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

log10_() → Any¶: 原地计算 TensorDict 中每个元素的 log10() 值。

log1p() → Any¶: 计算 TensorDict 每个元素的 log1p() 值。

log1p_() → Any¶: 原地计算 TensorDict 每个元素的 log1p() 值。

log2() → Any¶: 计算 TensorDict 每个元素的 log2() 值。

log2_() → Any¶: 原地计算 TensorDict 每个元素的 log2() 值。

log_() → Any¶: 原地计算 TensorDict 每个元素的 log() 值。

logical_and(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

对 self 和 other 执行逻辑 AND 操作。

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

参数:: other (TensorDictBase or torch.Tensor) – 用于执行逻辑 AND 操作的张量或 TensorDict。
关键字参数:: default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

logsumexp(dim=None, keepdim=False, *, out=None)¶

返回给定维度 dim 上输入 tensordict 各行的指数对数之和。计算是数值稳定的。

如果 keepdim 为 True，则输出张量的大小与输入张量相同，但在 dim 维度上大小为 1。否则，dim 将被压缩（参见 squeeze()），导致输出张量比输入张量少 1 个（或 len(dim) 个）维度。

参数:

dim (int or tuple of ints, optional) – 要约简的维度或维度。如果为 None，则约简 tensordict 的所有批次维度。
keepdim (bool) – 输出 tensordict 是否保留维度。

关键字参数:

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

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

根据形状和可选的 dtype 创建一个空的内存映射张量。

警告

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

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

参数:

key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，将引发异常。
shape (torch.Size or equivalent, torch.Tensor for nested tensors) – 要写入的张量的形状。

关键字参数:

dtype (torch.dtype, optional) – 新张量的数据类型。

返回:

一个新的内存映射张量。

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

根据存储、形状和可选的 dtype 创建一个空的内存映射张量。

警告

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

注意

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

参数:

key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，将引发异常。
storage (torch.UntypedStorage) – 用于新 MemoryMappedTensor 的存储。必须是物理内存存储。
shape (torch.Size or equivalent, torch.Tensor for nested tensors) – 要写入的张量的形状。

关键字参数:

dtype (torch.dtype, optional) – 新张量的数据类型。

返回:

一个具有给定存储的新内存映射张量。

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, 可选) – 如果为 False，则新张量将共享输入张量的元数据（如形状和数据类型），但内容将为空。默认为 True。

返回:

一个具有给定存储的新内存映射张量。

map(fn: Callable, dim: int = 0, num_workers: Optional[int] = None, chunksize: Optional[int] = None, num_chunks: Optional[int] = None, pool: Optional[Pool] = None, generator: Optional[Generator] = None, max_tasks_per_child: Optional[int] = None, worker_threads: int = 1, mp_start_method: Optional[str] = None)¶

将一个函数映射到tensordict在某个维度上的切片。

此方法将通过将 tensordict 切块成相等大小的 tensordict 并分派操作到所需的 Worker 数量来应用一个函数。

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

注意

此方法特别适用于处理存储在磁盘上的大型数据集（例如内存映射的 tensordict），其中块将是原始数据的零拷贝切片，可以几乎零成本地传递给进程。这使得处理非常大的数据集（例如 TB 级别）的成本非常低。

参数:

fn (callable) – 要应用于 tensordict 的函数。支持类似于 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的签名。
dim (int, 可选) – tensordict 将被分块的维度。
num_workers (int, optional) – 工作线程的数量。与 pool 互斥。如果未提供，工作线程数量将设置为可用 CPU 的数量。

关键字参数:

out (TensorDictBase, 可选) – 输出的可选容器。其在提供的 dim 上的批次大小必须与 self.ndim 匹配。如果它被共享或内存映射（is_shared() 或 is_memmap() 返回 True），它将在远程进程中被填充，避免数据向内传输。否则，来自 self 切片的将被发送到进程，在当前进程中收集，然后就地写入 out。
chunksize (int, 可选) – 每个数据块的大小。chunksize 为 0 将沿所需维度解绑 tensordict，并在函数应用后重新堆叠它，而 chunksize>0 将分割 tensordict 并对结果列表调用 torch.cat()。如果未提供，则块的数量等于工作进程的数量。对于非常大的 tensordict，这样大的块可能无法放入内存中进行操作，并且可能需要更多的块来使操作在实践中可行。此参数与 num_chunks 互斥。
num_chunks (int, 可选) – 将 tensordict 分割成的块数。如果未提供，则块的数量等于工作进程的数量。对于非常大的 tensordict，这样大的块可能无法放入内存中进行操作，并且可能需要更多的块来使操作在实践中可行。此参数与 chunksize 互斥。
pool (mp.Pool, optional) – 要用于执行作业的多进程池实例。如果未提供，将在 map 方法中创建一个池。
generator (torch.Generator, optional) –
用于播种的生成器。将从其生成一个基础种子，并且池中的每个工作进程将使用提供的种子（加上一个唯一的整数 0 到 num_workers）进行播种。如果没有提供生成器，将使用一个随机整数作为种子。要与未播种的工作进程一起使用，应单独创建一个池并将其直接传递给 map()。

注意

当提供低数值种子时应谨慎，因为这可能导致实验之间的自相关。例如：如果请求 8 个工作线程且种子为 4，则工作线程的种子范围将是 4 到 11。如果种子是 5，则工作线程的种子范围将是 5 到 12。这两个实验将有 7 个种子的重叠，这可能对结果产生意外影响。

注意

为工作线程设置种子的目的是让每个工作线程拥有不同的种子，而不是为了让 map 方法的调用结果可重现。换句话说，两个实验可能并且很可能返回不同的结果，因为无法知道哪个工作线程会执行哪个作业。但是，我们可以确保每个工作线程都有不同的种子，并且每个工作线程上的伪随机操作将是不相关的。
max_tasks_per_child (int, 可选) – 每个子进程拾取的最大任务数。默认为 None，即对任务数没有限制。
worker_threads (int, 可选) – 工作进程的线程数。默认为 1。
index_with_generator (bool, 可选) – 如果为 True，则 tensordict 的分割/分块将在查询期间进行，从而节省初始化时间。请注意，chunk() 和 split() 比索引（在生成器中使用）效率更高，因此初始化时间的处理时间增益可能会对总运行时间产生负面影响。默认为 False。
pbar (bool, 可选) – 如果为 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.]])

map_iter(fn: Callable[[TensorCollection], TensorCollection | 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) → Iterator[T]¶

沿一个维度迭代地将函数映射到tensordict的切片。

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

此方法将通过将tensordict分块为大小相等的tensordicts并将操作分派到所需数量的工作进程来应用于tensordict实例。它将逐个生成结果。

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

注意

此方法特别适用于处理存储在磁盘上的大型数据集（例如内存映射的 tensordict），其中块将是原始数据的零拷贝切片，可以几乎零成本地传递给进程。这使得处理非常大的数据集（例如 TB 级别）的成本非常低。

注意

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

参数:

fn (callable) – 要应用于 tensordict 的函数。支持类似于 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的签名。
dim (int, 可选) – tensordict 将被分块的维度。
num_workers (int, optional) – 工作线程的数量。与 pool 互斥。如果未提供，工作线程数量将设置为可用 CPU 的数量。

关键字参数:

shuffle (bool, 可选) – 索引是否应全局混洗。如果为 True，每个批次将包含不连续的样本。如果 index_with_generator=False 且 shuffle=True，将引发错误。默认为 False。
chunksize (int, 可选) – 每个数据块的大小。chunksize 为 0 将沿所需维度解绑 tensordict，并在函数应用后重新堆叠它，而 chunksize>0 将分割 tensordict 并对结果列表调用 torch.cat()。如果未提供，则块的数量等于工作进程的数量。对于非常大的 tensordict，这样大的块可能无法放入内存中进行操作，并且可能需要更多的块来使操作在实践中可行。此参数与 num_chunks 互斥。
num_chunks (int, 可选) – 将 tensordict 分割成的块数。如果未提供，则块的数量等于工作进程的数量。对于非常大的 tensordict，这样大的块可能无法放入内存中进行操作，并且可能需要更多的块来使操作在实践中可行。此参数与 chunksize 互斥。
pool (mp.Pool, optional) – 要用于执行作业的多进程池实例。如果未提供，将在 map 方法中创建一个池。
generator (torch.Generator, optional) –
用于播种的生成器。将从其生成一个基础种子，并且池中的每个工作进程将使用提供的种子（加上一个唯一的整数 0 到 num_workers）进行播种。如果没有提供生成器，将使用一个随机整数作为种子。要与未播种的工作进程一起使用，应单独创建一个池并将其直接传递给 map()。

注意

当提供低数值种子时应谨慎，因为这可能导致实验之间的自相关。例如：如果请求 8 个工作线程且种子为 4，则工作线程的种子范围将是 4 到 11。如果种子是 5，则工作线程的种子范围将是 5 到 12。这两个实验将有 7 个种子的重叠，这可能对结果产生意外影响。

注意

为工作线程设置种子的目的是让每个工作线程拥有不同的种子，而不是为了让 map 方法的调用结果可重现。换句话说，两个实验可能并且很可能返回不同的结果，因为无法知道哪个工作线程会执行哪个作业。但是，我们可以确保每个工作线程都有不同的种子，并且每个工作线程上的伪随机操作将是不相关的。
max_tasks_per_child (int, 可选) – 每个子进程拾取的最大任务数。默认为 None，即对任务数没有限制。
worker_threads (int, 可选) – 工作进程的线程数。默认为 1。
index_with_generator (bool, 可选) –
如果为 True，则 tensordict 的分割/分块将在查询期间进行，从而节省初始化时间。请注意，chunk() 和 split() 比索引（在生成器中使用）效率更高，因此初始化时间的处理时间增益可能会对总运行时间产生负面影响。默认为 True。

注意

index_with_generator 的默认值对于 map_iter 和 map 是不同的，前者假定将拆分后的 TensorDict 存储在内存中成本过高。
pbar (bool, 可选) – 如果为 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.]])

masked_fill(*args, **kwargs)¶

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.]])

masked_fill_(*args, **kwargs)¶

用期望值填充与掩码对应的项。

参数:

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.]])

masked_select(mask: Tensor) → Any¶

屏蔽 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.]])

max(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None, return_indices: bool = True) → Self | torch.Tensor¶

返回输入 tensordict 中所有元素的最大值。

参数:

dim (int, optional) – 如果为 None，则返回一个无维度的 tensordict，其中包含所有叶子的最大值（如果可以计算）。如果为整数，则仅当该维度与 tensordict 的形状兼容时，才会在指定的维度上调用 max。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。
return_argmins (bool, 可选) – 当 dim 参数被传递时，max() 返回一个包含值和索引的命名元组。此方法的 TensorDict 等价物是返回一个具有条目 "values" 和 "indices" 的 tensorclass，其内部结构相同。默认为 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.max(dim=0)
max(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.max()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.max(reduce=True)
tensor(3.2942)

maximum(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

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

参数:: other (TensorDict 或 Tensor) – 另一个输入tensordict或张量。
关键字参数:: default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

maximum_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: maximum_ (maximum() 的就地版本。

注意

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

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

尝试使 TensorDicts 密集堆叠，并在需要时回退到懒惰堆叠。

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

mean(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

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

参数:

dim (int, int 元组, str, 可选) – 如果为 None，则返回一个无维度的 tensordict，其中包含所有叶子的平均值（如果可计算）。如果为整数或整数元组，则当且仅当此维度与 tensordict 的形状兼容时，才会在指定的维度上调用 mean。目前只允许 “feature” 字符串。使用 dim=”feature” 将在所有特征维度上进行归约。如果 reduce=True，将返回一个形状与 TensorDict 的批次大小相同的张量。否则，将返回一个具有与 self 相同结构的 Tensordict，其中特征维度已归约。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.mean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.mean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.mean(reduce=True)
tensor(-0.0547)
>>> td.mean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.mean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.mean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

memmap(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) → Any¶

将所有张量写入内存映射的 Tensor 中，并放入新的 tensordict。

参数:

prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。
copy_existing (bool) – 如果为 False（默认值），并且 tensordict 中某项已是存储在磁盘上的张量且关联了文件，但未按 prefix 保存到正确位置，则会引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为 0。
return_early (bool, 可选) – 如果为 True 且 num_threads>0，则方法将返回一个 tensordict 的 future。
share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程之间共享，并且在单个节点内的任何工作进程上进行写入操作（例如就地更新或设置）将更新其他所有工作进程上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, 可选) – 如果为 False，则如果同一路径下已存在张量，将引发异常。默认为 True。

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

返回:: 返回一个新的 tensordict，其中张量存储在磁盘上（如果 return_early=False），否则返回一个 TensorDictFuture 实例。

注意

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

memmap_(prefix: Optional[str] = None, copy_existing: bool = False, num_threads: int = 0) → TensorDictBase¶

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

参数:

prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。
copy_existing (bool) – 如果为 False（默认值），并且 tensordict 中某项已是存储在磁盘上的张量且关联了文件，但未按 prefix 保存到正确位置，则会引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为 0。
return_early (bool, optional) – 如果为 True 且 num_threads>0，则该方法将返回一个 tensordict 的 future。可以使用 future.result() 查询结果 tensordict。
share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程之间共享，并且在单个节点内的任何工作进程上进行写入操作（例如就地更新或设置）将更新其他所有工作进程上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, 可选) – 如果为 False，则如果同一路径下已存在张量，将引发异常。默认为 True。

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

返回:: 如果 return_early=False，则返回 self，否则返回 TensorDictFuture 实例。

注意

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

memmap_like(prefix: str | None = None, copy_existing: bool = False, num_threads: int = 0) → Any¶

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

参数:

prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。
copy_existing (bool) – 如果为 False（默认值），并且 tensordict 中某项已是存储在磁盘上的张量且关联了文件，但未按 prefix 保存到正确位置，则会引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为 0。
return_early (bool, 可选) – 如果为 True 且 num_threads>0，则方法将返回一个 tensordict 的 future。
share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程之间共享，并且在单个节点内的任何工作进程上进行写入操作（例如就地更新或设置）将更新其他所有工作进程上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, 可选) – 如果为 False，则如果同一路径下已存在张量，将引发异常。默认为 True。

然后，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，则刷新其内容。

如果没有任何路径与之关联，此方法将引发异常。

min(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None, return_indices: bool = True) → Self | torch.Tensor¶

返回输入 tensordict 中所有元素的最小值。

参数:

dim (int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的最小值（如果可以计算）。如果为整数，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 min。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。
return_argmins (bool, 可选) – 当 dim 参数被传递时，min() 返回一个包含值和索引的命名元组。此方法的 TensorDict 等价物是返回一个具有条目 "values" 和 "indices" 的 tensorclass，其内部结构相同。默认为 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.min(dim=0)
min(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.min()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.min(reduce=True)
tensor(-2.9953)

minimum(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

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

参数:: other (TensorDict 或 Tensor) – 另一个输入tensordict或张量。
关键字参数:: default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

minimum_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: minimum_ (minimum() 的就地版本。

注意

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

modules() → Iterator[Module]¶

返回网络中所有模块的迭代器。

产生:: Module – 网络中的一个模块

注意

重复的模块只返回一次。在以下示例中，l 只返回一次。

示例

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device: Optional[Union[device, int]] = None) → Self¶

将所有模型参数和缓冲区移动到 MTIA。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 MTIA 上，则应在构建优化器之前调用它。

注意

此方法就地修改模块。

参数:: device (int, optional) – 如果指定，所有参数都将复制到该设备
返回:: self
返回类型:: 模块

mul(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

将 other 乘以 self。

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

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

参数:: other (TensorDict, Tensor 或 Number) – 从 self 中减去的张量或数字。
关键字参数:: default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

mul_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: mul_ (mul() 的就地版本。

注意

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

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

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

可调用签名必须是 Callable[Tuple[str, Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。

参数:

fn (Callable) – 要应用于 tensordict 中的（名称，张量）对的函数。对于每个叶子，只使用其叶子名称（不使用完整的 NestedKey）。
*others (TensorDictBase 实例, 可选) – 如果提供，这些 tensordict 实例应具有与 self 匹配的结构。fn 参数应接收与 tensordicts 数量（包括 self）相同的未命名输入。如果其他 tensordicts 存在缺失条目，可以通过 default 关键字参数传递默认值。
nested_keys (bool, optional) – 如果为 True，则将使用叶子的完整路径。默认为 False，即仅将最后一个字符串传递给函数。
batch_size (sequence of int, optional) – 如果提供，则结果 TensorDict 将具有所需的 batch_size。 batch_size 参数应与转换后的 batch_size 匹配。这是一个仅限关键字的参数。
device (torch.device, 可选) – 生成的设备，如果存在。
names (字符串列表, 可选) – 新的维度名称，以防batch_size被修改。
inplace (bool, optional) – 如果为 True，则就地进行更改。默认为 False。这是一个仅关键字参数。
default (Any, 可选) – 其他tensordict中缺失条目的默认值。如果未提供，缺失的条目将引发KeyError。
filter_empty (bool, optional) – 如果为 True，则将过滤掉空的 tensordicts。这还将降低计算成本，因为不会创建和销毁空的数据结构。为了向后兼容，默认为 False。
propagate_lock (bool, optional) – 如果为 True，则锁定的 tensordict 将产生另一个锁定的 tensordict。默认为 False。

call_on_nested (bool, optional) –

如果为 True，则该函数将应用于第一级张量和容器（TensorDict 或 tensorclass）。在这种情况下，func 负责传播其调用到嵌套级别。这允许在将调用传播到嵌套 tensordicts 时进行精细控制。如果为 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 执行的操作需要访问多个键来进行单次计算，则将 out 参数设置为 self 可能会导致操作产生静默错误的结果。例如

>>> 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)

named_buffers(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) → Iterator[tuple[str, torch.Tensor]]¶

返回模块缓冲区上的迭代器，同时生成缓冲区的名称和缓冲区本身。

参数:

prefix (str) – 要添加到所有缓冲区名称前的前缀。
recurse (bool) – 如果为 True，则会产生此模块及其所有子模块的缓冲区。否则，仅产生直接属于此模块的缓冲区。
remove_duplicate (bool, 可选) – 是否移除结果中的重复缓冲区。默认为 True。

产生:

(str, torch.Tensor) – 包含名称和缓冲区的元组

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

named_children() → Iterator[tuple[str, 'Module']]¶

返回对直接子模块的迭代器，生成模块的名称和模块本身。

产生:: (str, Module) – 包含名称和子模块的元组

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo: Optional[set['Module']] = None, prefix: str = '', remove_duplicate: bool = True)¶

返回网络中所有模块的迭代器，同时生成模块的名称和模块本身。

参数:

memo – 用于存储已添加到结果中的模块集合的 memo
prefix – 将添加到模块名称的名称前缀
remove_duplicate – 是否从结果中删除重复的模块实例

产生:

(str, Module) – 名称和模块的元组

注意

重复的模块只返回一次。在以下示例中，l 只返回一次。

示例

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) → Iterator[tuple[str, torch.nn.parameter.Parameter]]¶

返回模块参数的迭代器，同时生成参数的名称和参数本身。

参数:

prefix (str) – 要添加到所有参数名称前的前缀。
recurse (bool) – 如果为 True，则会产生此模块及其所有子模块的参数。否则，仅产生直接属于此模块的参数。
remove_duplicate (bool, 可选) – 是否移除结果中的重复参数。默认为 True。

产生:

(str, Parameter) – 包含名称和参数的元组

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

property names¶

tensordict 的维度名称。

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

另请参阅 refine_names() 以了解有关构造后设置名称的详细信息。

nanmean(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

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

参数:

dim (int, tuple of int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的平均值（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 mean。目前仅允许 “feature” 字符串。使用 dim=”feature” 将实现跨所有特征维度的约简。如果 reduce=True，则将返回一个形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个具有与 self 相同的结构但特征维度已约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.nanmean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.nanmean(reduce=True)
tensor(-0.0547)
>>> td.nanmean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.nanmean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

nansum(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回输入 tensordict 中所有非 NaN 元素的总和。

参数:

dim (int, tuple of int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的总和（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 sum。目前仅允许 “feature” 字符串。使用 dim=”feature” 将实现跨所有特征维度的约简。如果 reduce=True，则将返回一个形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个具有与 self 相同的结构但特征维度已约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nansum(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.nansum()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.nansum(reduce=True)
tensor(-0.)
>>> td.nansum(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nansum(reduce=True, dim="feature")
tensor([[15., 15., 15., 15.],
        [15., 15., 15., 15.],
        [15., 15., 15., 15.]])
>>> td.nansum(reduce=True, dim=0)
tensor([[9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.]])

property ndim: int¶: 参阅 batch_dims()。

ndimension() → int¶: 参阅 batch_dims()。

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

neg_() → Any¶: 原地计算 TensorDict 中每个元素的 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, empty_lazy: bool = False)¶

返回一个大小为 size 的 TensorDict，其中包含空的张量。

默认情况下，返回的 TensorDict 具有与此 tensordict 相同的 torch.dtype 和 torch.device。

参数:

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

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, optional) – 如果 autograd 应该记录返回张量的操作。默认为 False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。
pin_memory (bool, optional) – 如果设置为 True，则返回的张量将被分配到 pinned 内存中。仅适用于 CPU 张量。默认为 False。
empty_lazy (bool, optional) – 如果为 True，则 lazy stacks 将被清空其内容。当 lazy stack 的内容在填充新 tensordict 时可能会发生变化时，这会很有用。此参数将被传播到子 tensordict。默认为 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, empty_lazy: bool = False)¶

返回一个大小为 size 的 TensorDict，其中填充值为 1。

默认情况下，返回的 TensorDict 具有与此 tensordict 相同的 torch.dtype 和 torch.device。

参数:

size (sequence of int) – 定义输出张量形状的整数列表、元组或 torch.Size。
fill_value (scalar) – 用于填充输出 Tensor 的数值。

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, optional) – 如果 autograd 应该记录返回张量的操作。默认为 False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。
pin_memory (bool, optional) – 如果设置为 True，则返回的张量将被分配到 pinned 内存中。仅适用于 CPU 张量。默认为 False。
empty_lazy (bool, optional) – 如果为 True，则 lazy stacks 将被清空其内容。当 lazy stack 的内容在填充新 tensordict 时可能会发生变化时，这会很有用。此参数将被传播到子 tensordict。默认为 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, empty_lazy: bool = False)¶

返回一个大小为 size 的 TensorDict，其中填充值为 1。

默认情况下，返回的 TensorDict 具有与此 tensordict 相同的 torch.dtype 和 torch.device。

参数:

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

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, optional) – 如果 autograd 应该记录返回张量的操作。默认为 False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。
pin_memory (bool, optional) – 如果设置为 True，则返回的张量将被分配到 pinned 内存中。仅适用于 CPU 张量。默认为 False。
empty_lazy (bool, optional) – 如果为 True，则 lazy stacks 将被清空其内容。当 lazy stack 的内容在填充新 tensordict 时可能会发生变化时，这会很有用。此参数将被传播到子 tensordict。默认为 False。

new_tensor(data: torch.Tensor | tensordict._tensorcollection.TensorCollection, *, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, pin_memory: Optional[bool] = None) → Any¶

使用 data 作为张量返回一个新的 TensorDict。

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

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

参数:

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

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, optional) – 如果 autograd 应该记录返回张量的操作。默认为 False。
pin_memory (bool, optional) – 如果设置为 True，则返回的张量将被分配到 pinned 内存中。仅适用于 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, empty_lazy: bool = False)¶

返回一个大小为 size 的 TensorDict，其中填充值为 0。

默认情况下，返回的 TensorDict 具有与此 tensordict 相同的 torch.dtype 和 torch.device。

参数:

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

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, optional) – 如果 autograd 应该记录返回张量的操作。默认为 False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。
pin_memory (bool, optional) – 如果设置为 True，则返回的张量将被分配到 pinned 内存中。仅适用于 CPU 张量。默认为 False。
empty_lazy (bool, optional) – 如果为 True，则 lazy stacks 将被清空其内容。当 lazy stack 的内容在填充新 tensordict 时可能会发生变化时，这会很有用。此参数将被传播到子 tensordict。默认为 False。

non_tensor_items(include_nested: bool = False)¶: 递归地返回所有非张量叶节点。

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

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

关键字参数:

out (TensorDict, optional) – 输出 tensordict。
dtype (torch.dtype, optional) – 输出的 dtype (torch>=2.4)。

numel() → int¶

批次中的总元素数量。

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

numpy() → numpy.ndarray | dict[str, Any]¶

将 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!'}}

另请参阅 to_struct_array() 转换为结构化数组。

param_count(*, count_duplicates: bool = True) → int¶

计算参数数量（可索引项的总数），仅考虑张量。

关键字参数:: count_duplicates (bool) – 是否计算重复张量为独立张量。如果为 False，则仅计算严格相同的张量（来自同一基张量的不同视图但具有不同 id 的张量将被计算两次）。默认为 True（每个张量被假定为单个副本）。

parameters(recurse: bool = True) → Iterator[Parameter]¶

返回模块参数的迭代器。

这通常传递给优化器。

参数:: recurse (bool) – 如果为 True，则会产生此模块及其所有子模块的参数。否则，仅产生直接属于此模块的参数。
产生:: Parameter – 模块参数

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

permute(*args, **kwargs)¶

根据 dims 返回 tensordict 的视图，其中批次维度已重新排序。

参数:

*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(*args, **kwargs)¶

在存储的张量上调用 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) → Any¶

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

参数:: num_threads (int 或 str) – 如果提供，则用于在叶子上调用 pin_memory 的线程数。如果传递 "auto"，则自动确定线程数。

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

从 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

popitem()¶

移除最后插入 TensorDict 的项。

popitem 只会返回非嵌套值。

pow(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

将 self 的每个元素与 other 求幂，并返回结果张量。

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

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

参数:: other (float, tensor or tensordict) – 指数
关键字参数:: default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

pow_(other: tensordict.base.TensorDictBase | torch.Tensor) → Any¶: 原地版本的 pow()。

注意

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

prod(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回输入 TensorDict 中所有元素的乘积。

参数:

dim (int, tuple of int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的积（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 prod。目前仅允许 “feature” 字符串。使用 dim=”feature” 将实现跨所有特征维度的约简。如果 reduce=True，则将返回一个形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个具有与 self 相同的结构但特征维度已约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.prod(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.prod()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.prod(reduce=True)
tensor(-0.)
>>> td.prod(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.prod(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.prod(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

qint32() → Any¶: 将所有张量强制转换为 torch.qint32。

qint8() → Any¶: 将所有张量强制转换为 torch.qint8。

quint4x2() → Any¶: 将所有张量强制转换为 torch.quint4x2。

quint8() → Any¶: 将所有张量强制转换为 torch.quint8。

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

reciprocal_() → Any¶: 原地计算 TensorDict 中每个元素的 reciprocal() 值。

record_stream(stream: Stream) → Any¶

标记 TensorDict 已被此流使用。

当 TensorDict 被解除分配时，请确保在解除分配时该流上的所有工作完成之前，不要将张量内存重用于其他张量。

更多信息请参阅 record_stream()。

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

接收 TensorDict 的内容并用其更新内容。

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

参数:

src (int) – 源工作进程的 rank。

关键字参数:

group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `None`。
init_tag (int) – 用于标记张量的 init_tag。注意，这将根据 TensorDict 中包含的张量数量进行递增。
pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许从不同节点发送多个数据而不会重叠。请注意，这些伪随机数的生成成本很高（每秒 1e-5），这意味着它可能会减慢算法的运行速度。此值必须与传递给 send() 的值匹配。默认为 False。

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

在所有机器上归约 TensorDict。

只有具有 rank dst 的进程才能接收最终结果。

refine_names(*names) → Any¶

根据名称精炼 self 的维度名称。

精炼是重命名的特殊情况，它“提升”未命名维度。None 维度可以精炼为任何名称；已命名维度只能精炼为相同名称。

由于命名张量可以与未命名张量共存，因此细化名称提供了一种很好的方式来编写既适用于命名张量又适用于未命名张量的命名张量感知代码。

names 可以包含最多一个 Ellipsis (...)。Ellipsis 会贪婪地展开；它会就地展开以使 names 的长度与 self.dim() 相同，方法是使用 self.names 的相应索引处的名称。

Returns: 具有根据输入命名的维度的相同的 TensorDict。

示例

>>> 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"]

register_backward_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor], Union[tuple[torch.Tensor, ...], Tensor]], Union[None, tuple[torch.Tensor, ...], Tensor]]) → RemovableHandle¶

在模块上注册一个反向传播钩子。

此函数已弃用，建议使用 register_full_backward_hook()，并且此函数在未来版本中的行为将发生变化。

返回:: 一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子
返回类型:: torch.utils.hooks.RemovableHandle

register_buffer(name: str, tensor: Optional[Tensor], persistent: bool = True) → None¶

向模块添加一个缓冲区。

这通常用于注册一个不应被视为模型参数的 buffer。例如，BatchNorm 的 running_mean 不是一个参数，但它是 module 状态的一部分。Buffer 默认是持久化的，并且会与参数一起保存。通过将 persistent 设置为 False，可以改变这种行为。持久化 buffer 和非持久化 buffer 之间的唯一区别是，后者不会成为此 module 的 state_dict 的一部分。

可以使用给定名称作为属性访问缓冲区。

参数:

name (str) – 缓冲区的名称。可以使用给定名称从此模块访问缓冲区
tensor (Tensor or None) – 要注册的 buffer。如果为 None，则对 buffer 运行的操作（如 cuda）将被忽略。如果为 None，则 buffer **不**包含在 module 的 state_dict 中。
persistent (bool) – buffer 是否是此 module 的 state_dict 的一部分。

示例

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook(hook: Union[Callable[[T, tuple[Any, ...]], Any], Optional[Any]], Callable[[T, tuple[Any, ...], dict[str, Any], Any], Optional[Any]]], *, prepend: bool = False, with_kwargs: bool = False, always_call: bool = False) → RemovableHandle¶

在模块上注册一个前向钩子。

每次调用 forward() 计算出输出后，都会调用 hook。

如果 with_kwargs 为 False 或未指定，则输入仅包含传递给 module 的位置参数。关键字参数不会传递给 hook，只会传递给 forward。hook 可以修改输出。它可以就地修改输入，但这不会影响 forward，因为这是在 forward() 被调用之后才执行。hook 应该具有以下签名

hook(module, args, output) -> None or modified output

如果 with_kwargs 为 True，则前向钩子将接收传递给 forward 函数的 kwargs，并需要返回可能已修改的输出。钩子应该具有以下签名

hook(module, args, kwargs, output) -> None or modified output

参数:

hook (Callable) – 用户定义的待注册钩子。
prepend (bool) – 如果为 True，则提供的 hook 将在当前 torch.nn.Module 的所有现有 forward hook 之前触发。否则，提供的 hook 将在当前 torch.nn.Module 的所有现有 forward hook 之后触发。请注意，使用 register_module_forward_hook() 注册的全局 forward hook 将在为此方法注册的所有 hook 之前触发。默认值：False
with_kwargs (bool) – 如果为 True，则 hook 将接收传递给 forward 函数的 kwargs。默认值：False
always_call (bool) – 如果为 True，则无论在调用 Module 时是否引发异常，都会运行 hook。默认值：False

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook: Union[Callable[[T, tuple[Any, ...]], Optional[Any]], Callable[[T, tuple[Any, ...], dict[str, Any]], Optional[tuple[Any, dict[str, Any]]]], *, prepend: bool = False, with_kwargs: bool = False) → RemovableHandle¶

在模块上注册一个前向预钩子。

每次调用 forward() 之前，都会调用 hook。

如果 with_kwargs 为 false 或未指定，则输入仅包含传递给模块的位置参数。关键字参数不会传递给钩子，而只会传递给 forward。钩子可以修改输入。用户可以返回一个元组或单个修改后的值。我们将把值包装成一个元组，如果返回的是单个值（除非该值本身就是元组）。钩子应该具有以下签名

hook(module, args) -> None or modified input

如果 with_kwargs 为 true，则前向预钩子将接收传递给 forward 函数的 kwargs。如果钩子修改了输入，则应该返回 args 和 kwargs。钩子应该具有以下签名

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

参数:

hook (Callable) – 用户定义的待注册钩子。
prepend (bool) – 如果为 True，则提供的 hook 将在当前 torch.nn.Module 的所有现有 forward_pre hook 之前触发。否则，提供的 hook 将在当前 torch.nn.Module 的所有现有 forward_pre hook 之后触发。请注意，使用 register_module_forward_pre_hook() 注册的全局 forward_pre hook 将在为此方法注册的所有 hook 之前触发。默认值：False
with_kwargs (bool) – 如果为 True，则 hook 将接收传递给 forward 函数的 kwargs。默认值：False

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor], Union[tuple[torch.Tensor, ...], Tensor]], Union[None, tuple[torch.Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle¶

在模块上注册一个反向传播钩子。

每次计算相对于模块的梯度时，将调用此钩子，其触发规则如下：

通常，钩子在计算相对于模块输入的梯度时触发。

如果模块输入都不需要梯度，则在计算相对于模块输出的梯度时触发钩子。

如果模块输出都不需要梯度，则钩子将不触发。

钩子应具有以下签名

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

关于输入和输出的梯度（grad_input 和 grad_output）是包含输入的梯度和输出的梯度的元组。Hook 不应该修改其参数，但它可以选择性地返回一个关于输入的新的梯度，该梯度将被用于替换后续计算中的 grad_input。grad_input 只对应于作为位置参数给出的输入，并且所有关键字参数都会被忽略。对于所有非 Tensor 参数，grad_input 和 grad_output 中的条目将为 None。

由于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。类似地，调用者将接收模块前向函数返回的每个张量的视图。

警告

使用反向传播钩子时不允许就地修改输入或输出，否则将引发错误。

参数:

hook (Callable) – 要注册的用户定义钩子。
prepend (bool) – 如果为 True，则提供的 hook 将在当前 torch.nn.Module 的所有现有 backward hook 之前触发。否则，提供的 hook 将在当前 torch.nn.Module 的所有现有 backward hook 之后触发。请注意，使用 register_module_full_backward_hook() 注册的全局 backward hook 将在为此方法注册的所有 hook 之前触发。

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook: Callable[[Module, Union[tuple[torch.Tensor, ...], Tensor]], Union[None, tuple[torch.Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle¶

在模块上注册一个反向预钩子。

每次计算模块的梯度时，将调用此钩子。钩子应具有以下签名

hook(module, grad_output) -> tuple[Tensor] or None

grad_output 是一个元组。钩子不应修改其参数，但可以选择返回一个新的输出梯度，该梯度将取代 grad_output 用于后续计算。对于所有非 Tensor 参数，grad_output 中的条目将为 None。

由于技术原因，当此钩子应用于模块时，其前向函数将接收传递给模块的每个张量的视图。类似地，调用者将接收模块前向函数返回的每个张量的视图。

警告

使用反向传播钩子时不允许就地修改输入，否则将引发错误。

参数:

hook (Callable) – 要注册的用户定义钩子。
prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

返回:

一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子

返回类型:

torch.utils.hooks.RemovableHandle

register_get_post_hook(hook)¶: Register a hook to be called after any get operation on leaf tensors.

register_load_state_dict_post_hook(hook)¶

注册一个后钩子，用于在模块的 load_state_dict() 被调用后运行。

它应该具有以下签名：: hook(module, incompatible_keys) -> None

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

如果需要，可以就地修改给定的 incompatible_keys。

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

返回:: 一个句柄，可用于通过调用 handle.remove() 来移除添加的钩子
返回类型:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)¶

注册一个预钩子，用于在模块的 load_state_dict() 被调用之前运行。

它应该具有以下签名：: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

参数:: hook (Callable) – 在加载状态字典之前将调用的可调用钩子。

register_module(name: str, module: Optional[Module]) → None¶: Alias for add_module().

register_parameter(name: str, param: Optional[Parameter]) → None¶

向模块添加一个参数。

可以使用给定名称作为属性访问该参数。

参数:

name (str) – 参数的名称。可以使用给定名称从此模块访问该参数
param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

register_state_dict_post_hook(hook)¶

注册 state_dict() 方法的后置钩子。

它应该具有以下签名：: hook(module, state_dict, prefix, local_metadata) -> None

注册的钩子可以就地修改 state_dict。

register_state_dict_pre_hook(hook)¶

注册 state_dict() 方法的前置钩子。

它应该具有以下签名：: hook(module, prefix, keep_vars) -> None

注册的钩子可用于在进行 state_dict 调用之前执行预处理。

rename(*names, **rename_map)¶

返回一个克隆的 TensorDict，其中维度已重命名。

示例

>>> 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)¶

Same as rename(), but executes the renaming in-place.

示例

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

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

用新字符串重命名一个键，并返回具有更新后的键名的相同 TensorDict。

参数:

old_key (str 或 nested key) – 要重命名的键。
new_key (str 或 nested key) – 条目的新名称。
safe (bool, optional) – 如果为 True，则当新键已存在于 TensorDict 中时将引发错误。

返回:

self

repeat(*repeats: int) → Any¶

沿指定维度重复此张量。

Unlike expand(), this function copies the tensor’s data.

警告

repeat() behaves differently from repeat(), but is more similar to numpy.tile(). For the operator similar to numpy.repeat(), see repeat_interleave().

参数:: repeat (torch.Size, int..., tuple of int or list of int) – 沿每个维度重复此张量的次数。

示例

>>> import torch
>>>
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     {
...         "a": torch.randn(3, 4, 5),
...         "b": TensorDict({
...             "c": torch.randn(3, 4, 10, 1),
...             "a string": "a string!",
...         }, batch_size=[3, 4, 10])
...     }, batch_size=[3, 4],
... )
>>> print(td.repeat(1, 2))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 8, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                a string: NonTensorData(data=a string!, batch_size=torch.Size([3, 8, 10]), device=None),
                c: Tensor(shape=torch.Size([3, 8, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 8, 10]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 8]),
    device=None,
    is_shared=False)

repeat_interleave(*shape: int)¶

重复 TensorDict 的元素。

警告

这与 repeat() 不同，但与 numpy.repeat() 类似。

参数:

repeats (torch.Tensor or int) – 每个元素的重复次数。repeats 将被广播以适应给定轴的形状。
dim (int, optional) – 重复值的维度。默认为使用展平的输入数组，并返回一个展平的输出数组。

关键字参数:

output_size (int, optional) – Total output size for the given axis (e.g. sum of repeats). If given, it will avoid stream synchronization needed to calculate output shape of the tensordict.

返回:

重复的 TensorDict，其形状与输入相同，除了在给定轴上。

示例

>>> import torch
>>>
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     {
...         "a": torch.randn(3, 4, 5),
...         "b": TensorDict({
...             "c": torch.randn(3, 4, 10, 1),
...             "a string": "a string!",
...         }, batch_size=[3, 4, 10])
...     }, batch_size=[3, 4],
... )
>>> print(td.repeat_interleave(2, dim=0))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([6, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                a string: NonTensorData(data=a string!, batch_size=torch.Size([6, 4, 10]), device=None),
                c: Tensor(shape=torch.Size([6, 4, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([6, 4, 10]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([6, 4]),
    device=None,
    is_shared=False)

replace(*args, **kwargs)¶

创建 TensorDict 的浅拷贝，其中条目已被替换。

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

返回:: 如果输入非空，则返回 self 的副本，并更新条目。如果提供了空字典或未提供字典，并且 kwargs 为空，则返回 self。

requires_grad_(requires_grad=True) → Any¶

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

返回此 TensorDict。

参数:: requires_grad (bool, optional) – autograd 是否应该记录此 tensordict 上的操作。默认为 True。

reshape(*shape: int)¶

返回所需形状的连续、重塑的张量。

参数:: *shape (int) – 结果 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() → Any¶: 计算 TensorDict 中每个元素的 round() 值。

round_() → Any¶: 原地计算 TensorDict 中每个元素的 round() 值。

rsub(other: tensordict.base.TensorDictBase | torch.Tensor | float, *, alpha: float | None = None, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

从 self 中减去 other 乘以 alpha 的缩放值。

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

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

参数:

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

关键字参数:

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

save(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → Any¶

将tensordict保存到磁盘。

此函数是 memmap() 的代理。

property saved_path¶

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

此参数在 is_memmap() 返回 False（例如，当 TensorDict 解锁时）后即失效。

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

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

值不会被复制：对原始tensordict或新tensordict的张量的就地修改将导致两个tensordict都发生变化。

参数:

*keys (str) – 要选择的键。
inplace (bool) – 如果为 True，则 tensordict 将就地剪枝。默认为 False。
strict (bool, optional) – 选择一个不存在的键是否会返回错误。默认为 True。

返回:

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

注意

To select keys in a tensordict and return a version of this tensordict deprived of these keys, see the split_keys() method.

示例

>>> 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: 'torch.distributed.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → None¶

将 tensordict 的内容发送到远程工作进程。

参数:

dst (int) – 应将内容发送到的目标工作进程的 rank。

关键字参数:

group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `None`。
init_tag (int) – 用于标记张量的初始 init_tag。注意，此值将根据 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://: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://: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()

separates(*keys: NestedKey, default: Any = _NoDefault.ZERO, strict: bool = True, filter_empty: bool = True) → Any¶

将指定的键从 tensordict 中分离（就地操作）。

另请参阅

此方法等效于在单个分割上调用 split_keys() 且 inplace=True。

另请参阅

此方法等效于调用 exclude()，只是它返回数据的另一个分割。

参数:

keys (NestedKey) – 要从 tensordict 中分离的键。
default (Any, optional) – 当键丢失时返回的值。如果未指定且 strict=True，则会引发异常。否则，任何缺失键的默认值将是 None，除非另有指定。
strict (bool, optional) – 如果为 True，则在键丢失时引发异常。默认为 True。
filter_empty (bool, optional) – 如果为 True，则会删除 self 中的空 tensordicts。默认为 True。

返回:

分离出的 tensordict。

返回类型:

T

示例

>>> td = TensorDict(
...     a=0,
...     b=0,
...     c=0,
...     d=0,
... )
>>> td_a_c = td.separates("a", "c")
>>> print(td_a_c)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(td)
TensorDict(
    fields={
        b: 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)

set(key: NestedKey, item: Union[Tensor, TensorCollection], inplace: bool = False, **kwargs: Any) → TensorDictBase¶

设置一个新的键值对。

参数:

key (str, str 元组) – 要设置的键的名称。
item (torch.Tensor or equivalent, TensorDictBase instance) – 要存储在 tensordict 中的值。
inplace (bool, optional) – if True and if a key matches an existing key in the tensordict, then the update will occur in-place for that key-value pair. If inplace is True and the entry cannot be found, it will be added. For a more restrictive in-place operation, use set_() instead. Defaults to False.

关键字参数:

non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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: Union[Tensor, TensorCollection]) → Any¶

设置现有键的值，同时保留原始存储。

参数:

key (str) – 值的名称
item (torch.Tensor or compatible type, TensorDictBase) – 要存储在 tensordict 中的值。

关键字参数:

non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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: Union[Tensor, TensorCollection], index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]]) → Any¶

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

参数:

key (str, str 元组) – 要修改的键。
value (torch.Tensor) – 在 index 处设置的值。
index (int, tensor or tuple) – 写入值的索引。

关键字参数:

non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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_extra_state(state: Any) → None¶

设置加载的 state_dict 中包含的额外状态。

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

参数:: state (dict) – state_dict 中的额外状态

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)

set_submodule(target: str, module: Module, strict: bool = False) → None¶

如果存在，设置由 target 给定的子模块，否则抛出错误。

注意

如果 strict 设置为 False（默认），该方法将替换现有子模块或在父模块存在的情况下创建新子模块。如果 strict 设置为 True，该方法将仅尝试替换现有子模块，并在子模块不存在时引发错误。

例如，假设您有一个 nn.Module A，它看起来像这样

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(图示了一个 nn.Module A。 A 包含一个嵌套子模块 net_b，该子模块本身有两个子模块 net_c 和 linear。 net_c 随后又有一个子模块 conv。)

要用一个新的 Linear 子模块覆盖 Conv2d，可以调用 set_submodule("net_b.net_c.conv", nn.Linear(1, 1))，其中 strict 可以是 True 或 False。

要将一个新的 Conv2d 子模块添加到现有的 net_b 模块中，可以调用 set_submodule("net_b.conv", nn.Conv2d(1, 1, 1))。

在上面，如果设置 strict=True 并调用 set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True)，则会引发 AttributeError，因为 net_b 中不存在名为 conv 的子模块。

参数:

target – 要查找的子模块的完全限定字符串名称。（要指定完全限定字符串，请参阅上面的示例。）
module – 要设置子模块的对象。
strict – 如果为 False，该方法将替换现有子模块或创建新子模块（如果父模块存在）。如果为 True，则该方法只会尝试替换现有子模块，如果子模块不存在则抛出错误。

抛出:

ValueError – If the target string is empty or if module is not an instance of nn.Module.
AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

setdefault(key: NestedKey, default: Union[Tensor, TensorCollection, Any], inplace: bool = False) → Union[Tensor, TensorCollection]¶

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

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

参数:

key (str 或 nested key) – 值的名称。
default (torch.Tensor or compatible type, TensorDictBase) – 如果键不存在，则存储在 tensordict 中的值。

返回:

tensordict 中 key 的值。如果 key 之前未设置，则为 default。

示例

>>> 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。

share_memory() → Self¶: 请参阅 torch.Tensor.share_memory_()。

share_memory_(*args, **kwargs)¶

将所有张量放入共享内存。

TensorDict 将被锁定，这意味着任何非原地写入操作都会引发异常（例如，重命名、设置或删除条目）。反之，一旦 tensordict 被解锁，share_memory 属性将变为 False，因为进程间的身份不再有保证。

返回:: self

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

sigmoid_() → Any¶: 原地计算 TensorDict 中每个元素的 sigmoid() 值。

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

sign_() → Any¶: 原地计算 TensorDict 中每个元素的 sign() 值。

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

sin_() → Any¶: 原地计算 TensorDict 中每个元素的 sin() 值。

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

sinh_() → Any¶: 原地计算 TensorDict 中每个元素的 sinh() 值。

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

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

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

softmax(dim: int, dtype: Optional[dtype] = None)¶

将 softmax 函数应用于 tensordict 元素。

参数:

dim (int 或 int 元组) – 将在其上计算 softmax 的 tensordict 维度。
dtype (torch.dtype, optional) – 返回张量的期望数据类型。如果指定，则在执行操作之前，输入张量将被转换为 dtype。这有助于防止数据类型溢出。

property sorted_keys: list[tensordict._nestedkey.NestedKey]¶

按字母顺序返回键。

不支持其他参数。

如果 TensorDict 被锁定，键将被缓存，直到 TensorDict 解锁以获得更快的执行。

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

像 torch.split 一样，使用指定的尺寸在给定维度上分割 TensorDict 中的每个张量。

返回一个包含分割块视图的 TensorDict 实例列表。

参数:

split_size (int or 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, default: Any = _NoDefault.ZERO, strict: bool = True, reproduce_struct: bool = False) → Tuple[T, ...]¶

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

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

参数:

key_sets (键的序列 Dict[in_key, out_key] 或 键列表) – 各个分割。
inplace (bool, optional) – 如果为 True，则 self 中的键将被原地移除。默认为 False。
default (Any, 可选) – 当 key 丢失时要返回的值。如果未指定且 strict=True，则会引发异常。
strict (bool, optional) – 如果为 True，则在键丢失时引发异常。默认为 True。
reproduce_struct (bool, optional) – 如果为 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() → Any¶: 计算 self 的逐元素平方根。

sqrt_() → Any¶: In-place version of sqrt().

squeeze(*args, **kwargs)¶

压缩维度在 -self.batch_dims+1 和 self.batch_dims-1 之间的所有张量，并将它们返回到一个新的 tensordict 中。

参数:: dim (Optional[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: int = 0, *, out=None)¶

沿给定维度将 tensordicts 堆叠成一个单一的 tensordict。

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

stack_from_tensordict(dim: int = 0, *, sorted: Optional[Union[bool, List[NestedKey]]] = None, out: Optional[Tensor] = None) → Tensor¶

将 tensordict 的所有条目堆叠成一个张量。

参数:

dim (int, optional) – 条目应堆叠的维度。

关键字参数:

sorted (bool or list of NestedKeys) – if True, the entries will be stacked in alphabetical order. If False (default), the dict order will be used. Alternatively, a list of key names can be provided and the tensors will be stacked accordingly. This incurs some overhead as the list of keys will be checked against the list of leaf names in the tensordict.
out (torch.Tensor, optional) – 用于堆叠操作的可选目标张量。

stack_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → Any¶

将条目堆叠到新的条目中，并可能移除原始值。

参数:

keys (NestedKey 序列) – 要堆叠的条目。

关键字参数:

out_key (NestedKey) – 堆叠输入的新的键名。
keep_entries (bool, optional) – 如果为 False，则 keys 中的条目将被删除。默认为 False。
dim (int, optional) – the dimension along which the stack must occur. Defaults to 0.

返回: self

示例

>>> td = TensorDict(a=torch.zeros(()), b=torch.ones(()))
>>> td.stack_tensors("a", "b", out_key="c")
>>> assert "a" not in td
>>> assert (td["c"] == torch.tensor([0, 1])).all()

state_dict(*args, destination=None, prefix='', keep_vars=False, flatten=True)¶

从 tensordict 中生成 state_dict。

state_dict 的结构将保持嵌套，除非将 flatten 设置为 True。

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

参数:

destination (dict, optional) – 如果提供，tensordict 的状态将更新到 dict 中，并返回相同的对象。否则，将创建一个 OrderedDict 并返回。默认：None。
prefix (str, optional) – 添加到张量名称的前缀，用于组合 state_dict 中的键。默认：''。
keep_vars (bool, optional) – 默认情况下，state_dict 中返回的 torch.Tensor 项将从 autograd 中分离。如果设置为 True，则不会执行分离。默认：False。
flatten (bool, optional) – 是否应使用 "." 字符展平结构。默认为 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: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回输入 tensordict 中所有元素的标准差值。

参数:

dim (int, tuple of int, optional) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶子节点的总和值（如果可计算）。如果为整数或整数元组，则仅当该维度与 tensordict 形状兼容时，才会在指定的维度上调用 std。目前只允许 “feature” 字符串。使用 dim=”feature” 将在所有特征维度上进行归约。如果 reduce=True，将返回一个形状与 TensorDict 的 batch-size 相同的张量。否则，将返回一个结构与 self 相同的、具有归约特征维度的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

correction (int) – difference between the sample size and sample degrees of freedom. Defaults to Bessel’s correction, correction=1.
reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.std(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.std()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.std(reduce=True)
tensor(1.0006)
>>> td.std(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.std(reduce=True, dim="feature")
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
>>> td.std(reduce=True, dim=0)
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])

sub(other: tensordict.base.TensorDictBase | torch.Tensor | float, *, alpha: float | None = None, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

将 other 缩放 alpha 后，从 self 中减去。

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

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

参数:

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

关键字参数:

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

sub_(other: tensordict.base.TensorDictBase | torch.Tensor | float, alpha: float | None = None)¶: In-place version of sub().

注意

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

sum(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回输入 tensordict 中所有元素的总和值。

参数:

dim (int, tuple of int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的总和（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 sum。目前仅允许 “feature” 字符串。使用 dim=”feature” 将实现跨所有特征维度的约简。如果 reduce=True，则将返回一个形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个具有与 self 相同的结构但特征维度已约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.sum(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.sum()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.sum(reduce=True)
tensor(-0.)
>>> td.sum(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.sum(reduce=True, dim="feature")
tensor([[15., 15., 15., 15.],
        [15., 15., 15., 15.],
        [15., 15., 15., 15.]])
>>> td.sum(reduce=True, dim=0)
tensor([[9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.]])

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

tan_() → Any¶: 原地计算 TensorDict 中每个元素的 tan() 值。

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

tanh_() → Any¶: 原地计算 TensorDict 中每个元素的 tanh() 值。

tensor_split(indices_or_sections: int | list[int] | tuple[int, ...] | torch.Tensor, dim=0) → tuple[tensordict.base.TensorDictBase, ...]¶

将 TensorDict 沿着 dim 维度根据 indices_or_sections 指定的索引或节数分割成多个子 tensordict，所有这些子 tensordict 都是输入张量的视图。

参数:

indices_or_sections (int or List(int) or tuple(int) or 1D tensor of ints) – If indices_or_sections is an integer n or a zero dimensional long tensordict with value n, input is split into n sections along dimension dim. If input is divisible by n along dimension dim, each section will be of equal size, input.size(dim) / n. If input is not divisible by n, the sizes of the first int(input.size(dim) % n) sections will have size int(input.size(dim) / n) + 1, and the rest will have size int(input.size(dim) / n). If indices_or_sections is a list or tuple of ints, or a one-dimensional long tensor, then input is split along dimension dim at each of the indices in the list, tuple or tensor. For instance, indices_or_sections=[2, 3] and dim=0 would result in the tensors input[:2], input[2:3], and input[3:]. If indices_or_sections is a tensor, it must be a zero-dimensional or one-dimensional long tensor on the CPU.
dim (int, optional) – 要分割张量的维度。默认为 0

示例

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

to(*args, **kwargs) → TensorDictBase¶

将 TensorDictBase 子类映射到另一个设备、dtype 或另一个 TensorDictBase 子类（如果允许）。

不允许将张量转换为新的 dtype，因为 tensordict 不会绑定到只包含单一 tensor dtype。

参数:

device (torch.device, optional) – tensordict 的期望设备。
dtype (torch.dtype, optional) – tensordict 的期望浮点数或复数 dtype。
tensor (torch.Tensor, optional) – 此 TensorDict 中所有张量的期望 dtype 和设备都取自此张量。

关键字参数:

non_blocking (bool, optional) – whether the operations should be blocking.
memory_format (torch.memory_format, optional) – 此 tensordict 中 4D 参数和缓冲区所需的内存格式。
batch_size (torch.Size, optional) – 输出 tensordict 的结果 batch-size。
other (TensorDictBase, optional) –
TensorDict 实例，其 dtype 和 device 是该 TensorDict 中所有张量的目标 dtype 和 device。

注意

由于 TensorDictBase 实例没有 dtype，因此 dtype 是从示例叶子节点收集的。如果存在多个 dtype，则不会进行 dtype 转换。
non_blocking_pin (bool, optional) –
如果为 True，则张量将在发送到设备之前被 pinned。这将异步进行，但可以通过 num_threads 参数进行控制。

注意

tensordict.pin_memory().to("cuda") 通常比 tensordict.to("cuda", non_blocking_pin=True) 慢，因为在后一种情况下，pin_memory 是异步调用的。多线程 pin_memory 通常在张量大且数量多时很有益：当张量数量太少时，创建线程和收集数据的开销会超过多线程的优势；如果张量很小，遍历长列表的开销也可能过大。
num_threads (int or None, optional) – if non_blocking_pin=True, the number of threads to be used for pin_memory. By default, max(1, torch.get_num_threads()) threads will be spawn. num_threads=0 will cancel any multithreading for the pin_memory() calls.
inplace (bool, optional) – if True, the data will be written in-place in the same tensordict. This can be significantly faster whenever building a tensordict is CPU-overhead bound. Defaults to False.

返回:

如果设备与 tensordict 设备不同，以及/或传递了 dtype，则返回一个新的 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(*, retain_none: bool = True, convert_tensors: Union[bool, Literal['numpy']] = False, tolist_first: bool = False) → dict[str, Any]¶

返回一个键值对与 tensordict 匹配的字典。

参数:

retain_none (bool) – 如果为 True，则 tensorclass 实例中的 None 值将被写入字典。否则，它们将被丢弃。默认值：True。
convert_tensors (bool, "numpy") – 如果为 True，则在创建字典时，张量将被转换为列表。如果为“numpy”，则张量将被转换为 numpy 数组。否则，它们将保持为张量。默认值：False。
tolist_first (bool) – 如果为 True，则当 tensordict 具有批处理维度时，它将首先转换为列表。默认值：False。

返回:

tensordict 的字典表示。

另请参阅

tolist()

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     a=torch.arange(24).view(2, 3, 4),
...     b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)),
...     batch_size=(2, 3)
... )
>>> print(td.to_dict())
{'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]]]), ‘b’: {‘c’: tensor([[[ 0, 1], [ 2, 3], [ 4, 5]],

[[ 6, 7],
[ 8, 9], [10, 11]]])}}

>>> print(td.to_dict(convert_tensors=True))
{'a': [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]], 'b': {'c': [[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]}}

to_empty(*, device: Optional[Union[int, str, device]], recurse: bool = True) → Self¶

将参数和缓冲区移动到指定设备，而不复制存储。

参数:

device (torch.device) – The desired device of the parameters and buffers in this module. – 此模块中参数和缓冲区的目标设备。
recurse (bool) – 是否递归地将子模块的参数和缓冲区移动到指定的设备。

返回:

self

返回类型:

模块

to_h5(filename, **kwargs)¶

将 tensordict 转换为具有 h5 后端的 PersistentTensorDict。

参数:

filename (str 或 path) – h5 文件的路径。
**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_lazystack(dim: int = 0)¶

将 TensorDict 转换为 LazyStackedTensorDict 或等价物。

注意

此方法可用于更改 LazyStackedTensorDict 的堆叠维度。例如，如果您有一个 stack_dim=1 的 LazyStackedTensorDict，您可以使用此方法将其更改为 stack_dim=0。

>>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3))
>>> td2 = td.to_lazystack()
>>> td2.batch_size
torch.Size([2, 3])
>>> assert isinstance(td2, LazyStackedTensorDict)
>>> assert td2.stack_dim == 0
>>> td3 = td2.to_lazystack(1)
>>> assert td3.stack_dim == 1
>>> td3.batch_size
torch.Size([2, 3])

参数:: dim (int, optional) – 要堆叠 tensordict 的维度。默认为 0。
返回:: LazyStackedTensorDict 实例。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3))
>>> td2 = td.to_lazystack()
>>> td2.batch_size
torch.Size([2, 3])
>>> assert isinstance(td2, LazyStackedTensorDict)

将 TensorCollection 的内容写入流式数据集。

关键字参数:

out (str | Tuple[str, str]) –
用于保存分片文件的输出数据集目录。
1. 如果 out 是本地目录，则分片文件将本地保存。
2. 如果 out 是远程目录，则将创建一个本地临时目录来
  缓存分片文件，然后将分片文件上传到远程位置。完成后，在分片上传后删除临时目录。
3. 如果 out 是 (local_dir, remote_dir) 的元组，则分片文件将保存在
  local_dir 中，并上传到远程位置。
columns (dict[str, str]) – 可选的列字典。如果未提供，将自动从张量集合中推断。
keep_local (bool) – 如果数据集已上传，是否在上传后保留本地数据集目录。默认为 False。
compression (str, optional) – 可选的压缩或压缩级别。默认为 None。
hashes (List[str], optional) – 可选的要应用于分片文件的哈希算法列表。默认为 None。
size_limit (Union[int, str], optional) – 可选的分片大小限制，超过此限制将开始新的分片。如果为 None，则将所有内容放入一个分片。也可以指定字节的可读格式，例如 "100kb" 表示 100 千字节（100*1024）等等。默认为 1 << 26。如果
writer – (MDSWriter, optional): 要使用的写入器。将从 out 关键字参数和其他输入关键字参数创建。
**kwargs (Any) – 写入器的其他设置。

注意

MDSWriter 对嵌套字典的支持有限。处理嵌套 tensordict 的正确方法是先使用 flatten_keys() 方法，然后在读取后使用 unflatten_keys() 方法。

警告

此方法需要安装 mosaicml-streaming。

警告

对于非张量数据，数据类型必须是固定的。tensordict 恢复数据类型的方式是查看列表的第一个元素。如果为 None，将引发错误。否则，列表中的所有数据都必须具有相同的类型（或为 None 表示缺失）。

另请参阅

请参阅 Mosaic streaming 库 API https://docs.mosaicml.com/projects/streaming

以下示例展示了如何创建数据集并在 PyTorch dataloader 中加载它的端到端示例。

示例

>>> import tempfile
>>> from typing import Any
>>> from tensordict import TensorDict, LazyStackedTensorDict
>>> import torch
>>>
>>>
>>> td = LazyStackedTensorDict(
...     TensorDict(a=0, b=1, c=torch.randn(2), d="a string"),
...     TensorDict(a=0, b=1, c=torch.randn(3), d="another string"),
...     TensorDict(a=0, b=1, c=torch.randn(3), d="yet another string"),
... )
>>>
>>> with tempfile.TemporaryDirectory() as tempdir:
...     # Create a dataset on one process / thread / node...
...     td.to_mds(out=tempdir)
...
...     # Create a dataloader
...     from streaming import StreamingDataset
...
...     # Load the dataset on another thread / process / node...
...     dataset = StreamingDataset(local=tempdir, remote=None, batch_size=2)
...
...     # Use the class `from_list` method as a collate_fn
...     dl = torch.utils.data.DataLoader(dataset=dataset, batch_size=2, collate_fn=LazyStackedTensorDict.from_list)
...     for batch in dl:
...         print("batch", batch)

to_module(module: 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 属性。

to_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 hooks）。默认为 False。
non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

示例

>>> 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.data.zero_()
>>> params.to_module(module)
>>> assert (module.layers[0].linear1.weight == 0).all()

将 tensordict 用作上下文管理器可以方便地进行函数调用： .. rubric:: 示例

>>> from tensordict import from_module
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = TensorDict.from_module(module)
>>> params = params.data * 0 # Use TensorDictParams to remake these tensors regular nn.Parameter instances
>>> with params.to_module(module):
...     # Call the module with zeroed params
...     y = module(*inputs)
>>> # The module is repopulated with its original params
>>> assert (TensorDict.from_module(module) != 0).any()

返回:: 一个 tensordict，如果 return_swap 为 True，则包含来自模块的值，否则为 None。

to_namedtuple(dest_cls: Optional[type] = None) → Any¶

将 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: Optional[NestedKey] = None) → Any¶

将所有嵌套张量转换为填充版本并相应地调整批次大小。

参数:

padding (float) – tensordict 中张量的填充值。默认为 0.0。
mask_key (NestedKey, optional) – 如果提供，则为写入有效值掩码的键。如果异构维度不是 tensordict batch-size 的一部分，则会导致错误。默认为 None

to_pytree()¶

将 tensordict 转换为 PyTree。

如果 tensordict 不是从 pytree 创建的，此方法将仅返回 self 而不进行修改。

有关更多信息和示例，请参阅 from_pytree()。

to_struct_array() → ndarray¶

将 tensordict 转换为 numpy 结构化数组。

在 from_struct_array() - to_struct_array() 循环中，输入和输出数组的内容应匹配。但是，to_struct_array 不会保留原始数组的内存内容。

另请参阅

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

另请参阅

转换为 numpy 数组字典，请参阅 numpy()。

返回:: 输入 TensorDict 的 numpy 结构化数组表示。

示例

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({'a': torch.tensor([1, 2, 3]), 'b': torch.tensor([4.0, 5.0, 6.0])}, batch_size=[3])
>>> arr = td.to_struct_array()
>>> print(arr)
[(1, 4.) (2, 5.) (3, 6.)]

to_tensordict(*, retain_none: bool | None = None)¶

从 TensorDictBase 返回一个常规的 TensorDict 实例。

参数:: retain_none (bool) – 如果为 True，则 tensorclass 实例中的 None 值将被写入 tensordict。否则，它们将被丢弃。默认值：True。
返回:: 包含相同值的新 TensorDict 对象。

tolist(*, convert_nodes: bool = True, convert_tensors: Union[bool, Literal['numpy']] = False, tolist_first: bool = False, as_linked_list: bool = False) → List[Any]¶

返回 tensordict 的嵌套列表表示。

如果 tensordict 没有批次维度，则此方法返回单个列表或字典。否则，它返回一个嵌套列表，其中每个内部列表代表一个批次维度。

参数:

convert_nodes (bool) – 如果为 True，则叶子节点将被转换为字典。否则，它们将被作为值列表返回。默认值：True。
convert_tensors (bool, "numpy") – 如果为 True，则在创建字典时，张量将被转换为列表。如果为“numpy”，则张量将被转换为 numpy 数组。否则，它们将保持为张量。默认值：False。
tolist_first (bool) – 如果为 True，则当 tensordict 具有批处理维度时，它将首先转换为列表。默认值：False。
as_linked_list (bool) – 如果为 True，则列表将被转换为 tensordict.utils.LinkedList，它将在修改列表时自动更新 tensordict。默认值：False。

返回:

tensordict 的嵌套列表表示。

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     a=torch.arange(24).view(2, 3, 4),
...     b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)),
...     batch_size=(2, 3)
... )
>>> print(td.tolist(tolist_first=True))
[[{'a': tensor([0, 1, 2, 3]), 'b': [{'c': tensor(0)}, {'c': tensor(1)}]}, {'a': tensor([4, 5, 6, 7]), 'b': [{'c': tensor(2)}, {'c': tensor(3)}]}, {'a': tensor([ 8,  9, 10, 11]), 'b': [{'c': tensor(4)}, {'c': tensor(5)}]}], [{'a': tensor([12, 13, 14, 15]), 'b': [{'c': tensor(6)}, {'c': tensor(7)}]}, {'a': tensor([16, 17, 18, 19]), 'b': [{'c': tensor(8)}, {'c': tensor(9)}]}, {'a': tensor([20, 21, 22, 23]), 'b': [{'c': tensor(10)}, {'c': tensor(11)}]}]]
>>> print(td.tolist(tolist_first=False))
[[{'a': tensor([0, 1, 2, 3]), 'b': {'c': tensor([0, 1])}}, {'a': tensor([4, 5, 6, 7]), 'b': {'c': tensor([2, 3])}}, {'a': tensor([ 8,  9, 10, 11]), 'b': {'c': tensor([4, 5])}}], [{'a': tensor([12, 13, 14, 15]), 'b': {'c': tensor([6, 7])}}, {'a': tensor([16, 17, 18, 19]), 'b': {'c': tensor([8, 9])}}, {'a': tensor([20, 21, 22, 23]), 'b': {'c': tensor([10, 11])}}]]
>>> print(td.tolist(convert_tensors=False))
[[{'a': [0, 1, 2, 3], 'b': [{'c': 0}, {'c': 1}]}, {'a': [4, 5, 6, 7], 'b': [{'c': 2}, {'c': 3}]}, {'a': [8, 9, 10, 11], 'b': [{'c': 4}, {'c': 5}]}], [{'a': [12, 13, 14, 15], 'b': [{'c': 6}, {'c': 7}]}, {'a': [16, 17, 18, 19], 'b': [{'c': 8}, {'c': 9}]}, {'a': [20, 21, 22, 23], 'b': [{'c': 10}, {'c': 11}]}]]
>>> print(td.tolist(convert_nodes=False))
[[[tensor([0, 1, 2, 3]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([4, 5, 6, 7]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([ 8,  9, 10, 11]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)]], [[tensor([12, 13, 14, 15]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([16, 17, 18, 19]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([20, 21, 22, 23]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)]]]

train(mode: bool = True) → Self¶

将模块设置为训练模式。

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc. – 这只对某些模块有影响。有关其在训练/评估模式下的行为的详细信息，例如它们是否受影响，请参阅特定模块的文档，例如 Dropout、BatchNorm 等。

参数:: mode (bool) – 设置训练模式（True）或评估模式（False）。默认：True。
返回:: self
返回类型:: 模块

transpose(dim0, dim1)¶

返回输入张量的转置张量。给定的维度 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() → Any¶: 计算 TensorDict 每个元素的 trunc() 值。

trunc_() → Any¶: 原地计算 TensorDict 每个元素的 trunc() 值。

type(dst_type: dtype) → Any¶

将所有张量转换为 dst_type。

参数:: dst_type (type or string) – 所需类型

uint16() → Any¶: 将所有张量转换为 torch.uint16。

uint32() → Any¶: 将所有张量转换为 torch.uint32。

uint64() → Any¶: 将所有张量转换为 torch.uint64。

uint8() → Any¶: 将所有张量转换为 torch.uint8。

unbind(dim: int) → tuple[T, ...]¶

返回一个已解绑的 tensordict 索引元组，沿指定的维度。

示例

>>> 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) → TensorDictBase¶

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

TensorDict 类型将被丢失，结果将是简单的 TensorDict 实例。嵌套 tensordicts 的元数据将从根推断：数据树中的所有实例将共享相同的批次大小、维度名称和设备。

参数:

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() 方法在处理 state-dicts 时尤其有用，因为它们可以无缝地将扁平字典转换为模仿模型结构的数数据结构。

示例

>>> 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_() → Any¶

解锁 tensordict 以进行非原地操作。

可用作装饰器。

有关更多详细信息，请参阅 lock_()。

unsqueeze(*args, **kwargs)¶

为介于 -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, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]] | tensordict.base.TensorDictBase, clone: bool = False, inplace: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None, is_leaf: Optional[Callable[[Type], bool]] = None, update_batch_size: bool = False, ignore_lock: bool = False) → TensorDictBase¶

使用来自字典或其他 TensorDict 的值来更新 TensorDict。

警告

update 在 try/except 块中调用时会损坏数据。不要在这些块中使用此方法来尝试捕获和修补执行过程中发生的错误。

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。
clone (bool, 可选) – 在设置输入（ tensor）字典中的张量之前是否克隆它们。默认为 False。
inplace (bool, 可选) – 如果为 True 并且某个键与 tensordict 中的现有键匹配，则更新将对该键值对进行原地更新。如果找不到该条目，则会将其添加。默认为 False。

关键字参数:

keys_to_update (NestedKeys 序列, 可选) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update(*keys_to_update))。
non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.
is_leaf (Callable[[Type], bool], optional) –
一个指示对象类型是应被视为叶子节点并被交换，还是应被视为张量集合的可调用对象。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。
update_batch_size (bool, optional) –
如果为 True，则 update 将尝试更新目标（self）的 batch_size，如果它与源的 batch_size 不匹配。默认为 False。

注意

在 batch_size 不匹配的情况下，LazyStackTensorDict 实例将被清空其内容，并将使用源 tensordict 的副本重新填充容器。

注意

此参数假定 keys_to_update 为空，并且 inplace=False。如果目标（self）的键不是源键的子集，则会引发异常，因为 TensorDict 将无法推断如何处理额外的目标条目。
ignore_lock (bool, optional) – 如果为 True，则可以更新任何 tensordict，无论其锁定状态如何。默认为 False。

注意

当使用 N 个元素的 LazyStackedTensorDict 更新另一个 M 个元素的 LazyStackedTensorDict 时，沿堆栈维度，update 方法会将额外 tensordict 的副本附加到目标（self）的懒惰堆栈中。这使得用户可以依赖 update 来逐步增加懒惰堆栈。

返回:: 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: Union[dict[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], T], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) → T¶

使用来自字典或其他 TensorDict 的值原地更新 TensorDict。

与 update() 不同，此函数如果键对 self 未知，则会引发错误。

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。
clone (bool, 可选) – 在设置输入（ tensor）字典中的张量之前是否克隆它们。默认为 False。

关键字参数:

keys_to_update (NestedKeys 序列, 可选) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update_(data_src.select(*keys_to_update))。
non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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: Union[dict[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], T], idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) → T¶

使用来自字典或其他 TensorDict 的值，在指定的索引处原地更新 TensorDict。

与 TensorDict.update 不同，此函数将在键未知于 TensorDict 时抛出错误。

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。
idx (int, torch.Tensor, 可迭代对象, slice) – 更新应发生的 tensordict 的索引。
clone (bool, 可选) – 在设置输入（ tensor）字典中的张量之前是否克隆它们。默认为 False。

关键字参数:

keys_to_update (NestedKeys序列, 可选) – 如果提供，将仅更新 key_to_update 中的键列表。
non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) → Iterator[Union[Tensor, TensorCollection]]¶

返回一个表示 tensordict 值的生成器。

参数:

include_nested (bool, 可选) – 如果为 `True`，则返回嵌套值。默认为 `False`。
leaves_only (bool, 可选) – 如果为 `False`，则仅返回叶子节点。默认为 `False`。
is_leaf (callable, optional) –
一个作用于类类型的可调用对象，返回一个布尔值，指示该类是否应被视为叶子节点。

注意

is_leaf 的目的是不是阻止递归调用嵌套的 tensordicts，而是为过滤标记某些类型，当 `leaves_only=True` 时。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，嵌套 tensordict 的结构仍将被遍历。换句话说，`is_leaf` 不控制递归深度，而是提供了一种当 `leaves_only=True` 时从结果中过滤掉某些类型的方法。这意味着树中的一个节点既可以是叶节点，也可以是具有子节点的节点。实际上，`is_leaf` 的默认值不包含 tensordict 和 tensorclass 实例作为叶节点。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。

关键字参数:

sort (bool, optional) – 是否应该对键进行排序。对于嵌套键，键将根据其连接的名称进行排序（即，("a", "key") 将被计为 "a.key" 以进行排序）。请注意，当处理大型 tensordicts 时，排序可能会产生显著的开销。默认为 False。

var(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回输入 tensordict 中所有元素的方差值。

参数:

dim (int, int 元组, 可选) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶子的总和值（如果可以计算）。如果为整数或整数元组，则当此维度与 tensordict 形状兼容时，将在指定的维度上调用 var。当前只允许 “feature” 字符串。使用 dim=”feature” 将在所有特征维度上执行归约。如果 reduce=True，则将返回形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个结构与 self 相同的、特征维度已归约的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

correction (int) – difference between the sample size and sample degrees of freedom. Defaults to Bessel’s correction, correction=1.
reduce (bool, optional) – 如果为 True，则将在所有 TensorDict 值上执行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.var(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.var()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.var(reduce=True)
tensor(1.0006)
>>> td.var(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.var(reduce=True, dim="feature")
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
>>> td.var(reduce=True, dim=0)
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])

view(*shape: int, size: list | tuple | torch.Size | None = None, batch_size: torch.Size | None = None)¶

返回一个 tensordict，其中包含与 tensordict batch_size 兼容的新形状的张量视图。

或者，可以将 dtype 作为第一个未命名参数提供。在这种情况下，所有张量都将以相应的 dtype 进行视图处理。请注意，这假定新的形状与提供的 dtype 兼容。有关 dtype 视图的更多信息，请参阅 view()。

参数:

*shape (int) – 结果 TensorDict 的新形状。
dtype (torch.dtype) – 或者，用于表示张量内容的 dtype。
size – 可迭代对象

关键字参数:

batch_size (torch.Size, 可选) – 如果提供了 dtype，则可以使用此关键字参数重置 batch_size。如果 view 是使用形状调用的，则此参数无效。

返回:

具有所需 batch_size 的新 tensordict。

示例

>>> 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: Tensor, other: torch.Tensor | tensordict.base.TensorDictBase, *, out: tensordict.base.TensorDictBase | None = None, pad: int | bool = None, update_batch_size: bool = False)¶

根据 condition 从 self 或 other 中选择元素，并返回一个 TensorDict。

参数:

condition (BoolTensor) – 当 True (非零) 时，返回 self，否则返回 other。
other (TensorDictBase 或标量) – 当 other 是标量时，值为标量；或者，当 condition 为 False 时的值为选定的值。

关键字参数:

out (TensorDictBase, 可选) – 输出 TensorDictBase 实例。
pad (标量, 可选) – 如果提供，源或目标 tensordict 中缺失的键将写入为 torch.where(mask, self, pad) 或 torch.where(mask, pad, other)。默认为 None，即不允许缺失的键。
update_batch_size (bool, optional) – 如果为 True 且提供了 out，则输出的批次大小将更新为匹配条件的批次大小。默认为 False。

xpu(device: Optional[Union[device, int]] = None) → Self¶

将所有模型参数和缓冲区移动到 XPU。

这也会使关联的参数和缓冲区成为不同的对象。因此，如果模块在优化时将驻留在 XPU 上，则应在构建优化器之前调用它。

注意

此方法就地修改模块。

参数:: device (int, optional) – 如果指定，所有参数都将复制到该设备
返回:: self
返回类型:: 模块

zero_() → Any¶: 原地将 tensordict 中的所有张量清零。

zero_grad(set_to_none: bool = True) → Any¶

递归地将 TensorDict 的所有梯度清零。

参数:: set_to_none (bool, 可选) – 如果为 True，则 tensor.grad 将为 None，否则为 0。默认为 True。

TensorDictParams¶

文档

教程

资源