tensordict.nn.dispatch

class tensordict.nn.dispatch(separator='_', source='in_keys', dest='out_keys', auto_batch_size: bool = True)

允许使用 kwargs 调用期望 TensorDict 的函数。

dispatch() 必须用在具有 in_keys（或 source 关键字参数指定的其他键源）和 out_keys（或另一个 dest 键列表）属性的模块中，这些属性指示要从 tensordict 读取和写入的键。包装的函数也应该有一个 tensordict 前导参数。

生成的函数将返回一个单独的张量（如果 out_keys 中只有一个元素），否则它将返回一个按模块的 out_keys 排序的元组。

dispatch() 可以用作方法，也可以用作类，当需要传递额外的参数时。

参数:

• separator (str, optional) – 当 in_keys 是字符串元组时，用于组合子键的分隔符。默认为 "_"。

• source (str or list of keys, optional) – 如果提供字符串，则指向包含要使用的输入键列表的模块属性。如果提供列表，则包含用作模块输入的键。默认为 "in_keys"，即 TensorDictModule 输入键列表的属性名称。

• dest (str or list of keys, optional) – 如果提供字符串，则指向包含要使用的输出键列表的模块属性。如果提供列表，则包含用作模块输出的键。默认为 "out_keys"，即 TensorDictModule 输出键列表的属性名称。

• auto_batch_size (bool, optional) – 如果为 True，则输入 tensordict 的 batch-size 将自动确定为所有输入张量之间公共维度的最大数量。默认为 True。

示例

>>> class MyModule(nn.Module):
...     in_keys = ["a"]
...     out_keys = ["b"]
...
...     @dispatch
...     def forward(self, tensordict):
...         tensordict['b'] = tensordict['a'] + 1
...         return tensordict
...
>>> module = MyModule()
>>> b = module(a=torch.zeros(1, 2))
>>> assert (b == 1).all()
>>> # equivalently
>>> class MyModule(nn.Module):
...     keys_in = ["a"]
...     keys_out = ["b"]
...
...     @dispatch(source="keys_in", dest="keys_out")
...     def forward(self, tensordict):
...         tensordict['b'] = tensordict['a'] + 1
...         return tensordict
...
>>> module = MyModule()
>>> b = module(a=torch.zeros(1, 2))
>>> assert (b == 1).all()
>>> # or this
>>> class MyModule(nn.Module):
...     @dispatch(source=["a"], dest=["b"])
...     def forward(self, tensordict):
...         tensordict['b'] = tensordict['a'] + 1
...         return tensordict
...
>>> module = MyModule()
>>> b = module(a=torch.zeros(1, 2))
>>> assert (b == 1).all()

dispatch_kwargs() 也可以使用默认的 "_" 分隔符处理嵌套键。

示例

>>> class MyModuleNest(nn.Module):
...     in_keys = [("a", "c")]
...     out_keys = ["b"]
...
...     @dispatch
...     def forward(self, tensordict):
...         tensordict['b'] = tensordict['a', 'c'] + 1
...         return tensordict
...
>>> module = MyModuleNest()
>>> b, = module(a_c=torch.zeros(1, 2))
>>> assert (b == 1).all()

如果想要其他分隔符，可以在构造函数中使用 separator 参数进行指定

示例

>>> class MyModuleNest(nn.Module):
...     in_keys = [("a", "c")]
...     out_keys = ["b"]
...
...     @dispatch(separator="sep")
...     def forward(self, tensordict):
...         tensordict['b'] = tensordict['a', 'c'] + 1
...         return tensordict
...
>>> module = MyModuleNest()
>>> b, = module(asepc=torch.zeros(1, 2))
>>> assert (b == 1).all()

由于输入键是排序的字符串序列，dispatch() 也可以与无名参数一起使用，其中顺序必须与输入键的顺序匹配。

注意

如果第一个参数是 TensorDictBase 实例，则假定 **不** 使用 dispatch，并且该 tensordict 包含运行模块所需的所有信息。换句话说，不能用指向 tensordict 实例的模块输入第一个键来分解 tensordict。通常，首选使用 dispatch() 处理仅叶节点 tensordict。

示例

>>> class MyModuleNest(nn.Module):
...     in_keys = [("a", "c"), "d"]
...     out_keys = ["b"]
...
...     @dispatch
...     def forward(self, tensordict):
...         tensordict['b'] = tensordict['a', 'c'] + tensordict["d"]
...         return tensordict
...
>>> module = MyModuleNest()
>>> b, = module(torch.zeros(1, 2), d=torch.ones(1, 2))  # works
>>> assert (b == 1).all()
>>> b, = module(torch.zeros(1, 2), torch.ones(1, 2))  # works
>>> assert (b == 1).all()
>>> try:
...     b, = module(torch.zeros(1, 2), a_c=torch.ones(1, 2))  # fails
... except:
...     print("oopsy!")
...