TensorDictSequential¶
- class tensordict.nn.TensorDictSequential(*args, **kwargs)¶
TensorDictModule 的序列。
类似于
nn.Sequence，它将一个张量通过一个读取和写入单个张量的映射链，这个模块将通过查询输入的每个模块来读写 tensordict。在调用TensorDictSequencial实例时，会期望参数列表（和缓冲区）被连接到单个列表中。- 参数:
modules (OrderedDict[str, Callable[[TensorDictBase], TensorDictBase]] | List[Callable[[TensorDictBase], TensorDictBase]]) – 按顺序排列的可调用对象，它们接受一个 TensorDictBase 作为输入并返回一个 TensorDictBase。这些可以是 TensorDictModuleBase 的实例，也可以是匹配此签名的任何其他函数。请注意，如果使用非 TensorDictModuleBase 可调用对象，其输入和输出键将不会被跟踪，因此不会影响 TensorDictSequential 的 in_keys 和 out_keys 属性。常规的
dict输入将在必要时转换为OrderedDict。- 关键字参数:
partial_tolerant (bool, optional) – 如果为 True，则输入 tensordict 可能缺少某些输入键。如果是这种情况，将只执行那些可以根据存在的键执行的模块。此外，如果输入 tensordict 是一个懒惰堆叠的 tensordicts 并且 partial_tolerant 为
True并且堆叠没有必需的键，那么 TensorDictSequential 将会扫描子 tensordicts 来查找具有必需键的 tensordicts（如果有）。默认为 False。selected_out_keys (iterable of NestedKeys, optional) – 要选择的 out-keys 列表。如果未提供，则将写入所有
out_keys。inplace (bool or str, optional) – 如果是 True，则输入 tensordict 被就地修改。如果是 False，则会创建一个新的空的
TensorDict实例。如果是 “empty”，则使用 input.empty()（即输出保留类型、设备和批次大小）。默认为 None（依赖于子模块）。
注意
一个
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 编码。
示例
>>> 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。
返回: 相同的模块，以就地修改方式返回，并更新了
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_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'])
此方法会传播到嵌套的序列
>>> 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'])
inplace 参数允许对输出类型进行精细控制，例如允许将计算图的结果写入输入对象，而不跟踪中间张量。
示例
>>> import torch >>> from tensordict import TensorClass >>> from tensordict.nn import TensorDictModule as Mod, TensorDictSequential as Seq >>> >>> class MyClass(TensorClass): ... input: torch.Tensor ... output: torch.Tensor | None = None >>> >>> obj = MyClass(torch.randn(2, 3), batch_size=(2,)) >>> >>> model = Seq( ... Mod( ... lambda x: (x + 1, x - 1), ... in_keys=["input"], ... out_keys=[("intermediate", "0"), ("intermediate", "1")], ... inplace=False ... ), ... Mod( ... lambda y0, y1: y0 * y1, ... in_keys=[("intermediate", "0"), ("intermediate", "1")], ... out_keys=["output"], ... inplace=False ... ), ... inplace=True, ) >>> print(model(obj)) MyClass( input=Tensor(shape=torch.Size([2, 3]), device=cpu, dtype=torch.float32, is_shared=False), output=Tensor(shape=torch.Size([2, 3]), device=cpu, dtype=torch.float32, is_shared=False), output=None, batch_size=torch.Size([2]), device=None, is_shared=False)
