目录
快捷方式

TensorDictSequential

class tensordict.nn.TensorDictSequential(*args, **kwargs)

TensorDictModule 的序列。

类似于 nn.Sequence，它将张量传递通过一系列映射，每个映射读取和写入单个张量，此模块将通过查询每个输入模块来读取和写入 tensordict。当使用函数式模块调用 TensorDictSequencial 实例时，预计参数列表（和缓冲区）将连接在单个列表中。

参数:

modules (TensorDictModules 的可迭代对象) – 要顺序运行的 TensorDictModule 实例的有序序列。

关键字参数:

  • partial_tolerant (bool, 可选) – 如果为 True，则输入 tensordict 可以缺少某些输入键。 如果是这样，则只会执行那些可以在给定存在的键的情况下执行的模块。 此外，如果输入 tensordict 是 tensordict 的惰性堆栈，并且 partial_tolerant 为 True，并且如果堆栈没有所需的键，则 TensorDictSequential 将扫描子 tensordict，查找任何具有所需键的子 tensordict。

  • selected_out_keys (NestedKeys 的可迭代对象, 可选) – 要选择的输出键列表。 如果未提供，将写入所有 out_keys。

注意

TensorDictSequential 实例可能具有很长的输出键列表，并且人们可能希望在执行后删除其中一些键以提高清晰度或节省内存。 如果是这种情况，可以在实例化后使用方法 select_out_keys()，或者可以将 selected_out_keys 传递给构造函数。

示例

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule, TensorDictSequential
>>> torch.manual_seed(0)
>>> module = TensorDictSequential(
...     TensorDictModule(lambda x: x+1, in_keys=["x"], out_keys=["x+1"]),
...     TensorDictModule(nn.Linear(3, 4), in_keys=["x+1"], out_keys=["w*(x+1)+b"]),
... )
>>> # with tensordict input
>>> print(module(TensorDict({"x": torch.zeros(3)}, [])))
TensorDict(
    fields={
        w*(x+1)+b: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
        x+1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        x: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> # with tensor input: returns all the output keys in the order of the modules, ie "x+1" and "w*(x+1)+b"
>>> module(x=torch.zeros(3))
(tensor([1., 1., 1.]), tensor([-0.7214, -0.8748,  0.1571, -0.1138], grad_fn=<AddBackward0>))
>>> module(torch.zeros(3))
(tensor([1., 1., 1.]), tensor([-0.7214, -0.8748,  0.1571, -0.1138], grad_fn=<AddBackward0>))

TensorDictSequence 支持函数式、模块化和 vmap 编码： .. rubric:: 示例

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import (
...     ProbabilisticTensorDictModule,
...     ProbabilisticTensorDictSequential,
...     TensorDictModule,
...     TensorDictSequential,
... )
>>> from tensordict.nn.distributions import NormalParamExtractor
>>> from tensordict.nn.functional_modules import make_functional
>>> from torch.distributions import Normal
>>> td = TensorDict({"input": torch.randn(3, 4)}, [3,])
>>> net1 = torch.nn.Linear(4, 8)
>>> module1 = TensorDictModule(net1, in_keys=["input"], out_keys=["params"])
>>> normal_params = TensorDictModule(
...      NormalParamExtractor(), in_keys=["params"], out_keys=["loc", "scale"]
...  )
>>> td_module1 = ProbabilisticTensorDictSequential(
...     module1,
...     normal_params,
...     ProbabilisticTensorDictModule(
...         in_keys=["loc", "scale"],
...         out_keys=["hidden"],
...         distribution_class=Normal,
...         return_log_prob=True,
...     )
... )
>>> module2 = torch.nn.Linear(4, 8)
>>> td_module2 = TensorDictModule(
...    module=module2, in_keys=["hidden"], out_keys=["output"]
... )
>>> td_module = TensorDictSequential(td_module1, td_module2)
>>> params = TensorDict.from_module(td_module)
>>> with params.to_module(td_module):
...     _ = td_module(td)
>>> print(td)
TensorDict(
    fields={
        hidden: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        input: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        loc: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        output: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        params: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        sample_log_prob: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
在 vmap 情况下
>>> from torch import vmap
>>> params = params.expand(4)
>>> def func(td, params):
...     with params.to_module(td_module):
...         return td_module(td)
>>> td_vmap = vmap(func, (None, 0))(td, params)
>>> print(td_vmap)
TensorDict(
    fields={
        hidden: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        input: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        loc: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        output: Tensor(shape=torch.Size([4, 3, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        params: Tensor(shape=torch.Size([4, 3, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        sample_log_prob: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=None,
    is_shared=False)
forward(tensordict: TensorDictBase = None, tensordict_out: tensordict.base.TensorDictBase | None >= None, **kwargs: Any) → TensorDictBase

当未设置 tensordict 参数时，kwargs 用于创建 TensorDict 的实例。

reset_out_keys()

将 out_keys 属性重置为其原始值。

返回: 相同的模块，具有其原始 out_keys 值。

示例

>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule, TensorDictSequential
>>> import torch
>>> mod = TensorDictModule(lambda x, y: (x+2, y+2), in_keys=["a", "b"], out_keys=["c", "d"])
>>> mod.select_out_keys("d")
>>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, [])
>>> mod(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: 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)
>>> mod.reset_out_keys()
>>> mod(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        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)
select_out_keys(*selected_out_keys) → TensorDictSequential

选择将在输出 tensordict 中找到的键。

当想要去除复杂图中的中间键，或者当这些键的存在可能触发意外行为时，这很有用。

原始 out_keys 仍然可以通过 module.out_keys_source 访问。

参数:

*out_keys (字符串或字符串元组的序列) – 应在输出 tensordict 中找到的输出键。

返回: 相同的模块，已就地修改，并更新了 out_keys。

最简单的用法是使用 TensorDictModule

示例

>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule, TensorDictSequential
>>> import torch
>>> mod = TensorDictModule(lambda x, y: (x+2, y+2), in_keys=["a", "b"], out_keys=["c", "d"])
>>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, [])
>>> mod(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        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)
>>> mod.select_out_keys("d")
>>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, [])
>>> mod(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: 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)

此功能也适用于分派的参数： .. rubric:: 示例

>>> mod(torch.zeros(()), torch.ones(()))
tensor(2.)

此更改将就地发生（即，将返回相同的模块，并更新 out_keys 列表）。 可以使用 TensorDictModuleBase.reset_out_keys() 方法还原。

示例

>>> mod.reset_out_keys()
>>> mod(TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        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)

这也适用于其他类，例如 Sequential： .. rubric:: 示例

>>> from tensordict.nn import TensorDictSequential
>>> seq = TensorDictSequential(
...     TensorDictModule(lambda x: x+1, in_keys=["x"], out_keys=["y"]),
...     TensorDictModule(lambda x: x+1, in_keys=["y"], out_keys=["z"]),
... )
>>> td = TensorDict({"x": torch.zeros(())}, [])
>>> seq(td)
TensorDict(
    fields={
        x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        y: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        z: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> seq.select_out_keys("z")
>>> td = TensorDict({"x": torch.zeros(())}, [])
>>> seq(td)
TensorDict(
    fields={
        x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        z: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
select_subsequence(in_keys: Optional[Iterable>[NestedKey]] = None, out_keys: Optional[Iterable>[NestedKey]] = None) → TensorDictSequential

返回一个新的 TensorDictSequential，其中仅包含计算给定输入键的给定输出键所需的模块。

参数:

  • in_keys – 我们要选择的子序列的输入键。 所有 in_keys 中缺少的键将被视为不相关，并且仅将这些键作为输入的模块将被丢弃。 生成的顺序模块将遵循模式“所有输出将受到任何 in <in_keys> 的不同值影响的模块”。 如果未提供，则假定为模块的 in_keys。

  • out_keys – 我们要选择的子序列的输出键。 只有获取 out_keys 所需的模块才会在生成的序列中找到。 生成的顺序模块将遵循模式“所有调节 <out_keys> 条目的值或值的模块。” 如果未提供，则假定为模块的 out_keys。

返回:

一个新的 TensorDictSequential，其中仅包含根据给定输入和输出键所需的模块。

示例

>>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod
>>> idn = lambda x: x
>>> module = Seq(
...     Mod(idn, in_keys=["a"], out_keys=["b"]),
...     Mod(idn, in_keys=["b"], out_keys=["c"]),
...     Mod(idn, in_keys=["c"], out_keys=["d"]),
...     Mod(idn, in_keys=["a"], out_keys=["e"]),
... )
>>> # select all modules whose output depend on "a"
>>> module.select_subsequence(in_keys=["a"])
TensorDictSequential(
    module=ModuleList(
      (0): TensorDictModule(
          module=<function <lambda> at 0x126ed1ca0>,
          device=cpu,
          in_keys=['a'],
          out_keys=['b'])
      (1): TensorDictModule(
          module=<function <lambda> at 0x126ed1ca0>,
          device=cpu,
          in_keys=['b'],
          out_keys=['c'])
      (2): TensorDictModule(
          module=<function <lambda> at 0x126ed1ca0>,
          device=cpu,
          in_keys=['c'],
          out_keys=['d'])
      (3): TensorDictModule(
          module=<function <lambda> at 0x126ed1ca0>,
          device=cpu,
          in_keys=['a'],
          out_keys=['e'])
    ),
    device=cpu,
    in_keys=['a'],
    out_keys=['b', 'c', 'd', 'e'])
>>> # select all modules whose output depend on "c"
>>> module.select_subsequence(in_keys=["c"])
TensorDictSequential(
    module=ModuleList(
      (0): TensorDictModule(
          module=<function <lambda> at 0x126ed1ca0>,
          device=cpu,
          in_keys=['c'],
          out_keys=['d'])
    ),
    device=cpu,
    in_keys=['c'],
    out_keys=['d'])
>>> # select all modules that affect the value of "c"
>>> module.select_subsequence(out_keys=["c"])
TensorDictSequential(
    module=ModuleList(
      (0): TensorDictModule(
          module=<function <lambda> at 0x126ed1ca0>,
          device=cpu,
          in_keys=['a'],
          out_keys=['b'])
      (1): TensorDictModule(
          module=<function <lambda> at 0x126ed1ca0>,
          device=cpu,
          in_keys=['b'],
          out_keys=['c'])
    ),
    device=cpu,
    in_keys=['a'],
    out_keys=['b', 'c'])
>>> # select all modules that affect the value of "e"
>>> module.select_subsequence(out_keys=["e"])
TensorDictSequential(
    module=ModuleList(
      (0): TensorDictModule(
          module=<function <lambda> at 0x126ed1ca0>,
          device=cpu,
          in_keys=['a'],
          out_keys=['e'])
    ),
    device=cpu,
    in_keys=['a'],
    out_keys=['e'])

此方法传播到嵌套的 sequential

>>> module = Seq(
...     Seq(
...         Mod(idn, in_keys=["a"], out_keys=["b"]),
...         Mod(idn, in_keys=["b"], out_keys=["c"]),
...     ),
...     Seq(
...         Mod(idn, in_keys=["b"], out_keys=["d"]),
...         Mod(idn, in_keys=["d"], out_keys=["e"]),
...     ),
... )
>>> # select submodules whose output will be affected by a change in "b" or "d" AND which output is "e"
>>> module.select_subsequence(in_keys=["b", "d"], out_keys=["e"])
TensorDictSequential(
    module=ModuleList(
      (0): TensorDictSequential(
          module=ModuleList(
            (0): TensorDictModule(
                module=<function <lambda> at 0x129efae50>,
                device=cpu,
                in_keys=['b'],
                out_keys=['d'])
            (1): TensorDictModule(
                module=<function <lambda> at 0x129efae50>,
                device=cpu,
                in_keys=['d'],
                out_keys=['e'])
          ),
          device=cpu,
          in_keys=['b'],
          out_keys=['d', 'e'])
    ),
    device=cpu,
    in_keys=['b'],
    out_keys=['d', 'e'])

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源

为了分析流量和优化您的体验，我们在本网站上使用 Cookie。 通过点击或导航，您同意允许我们使用 Cookie。 作为本网站的当前维护者，Facebook 的 Cookie 政策适用。 了解更多信息，包括有关可用控件的信息：Cookie 政策。