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, 可选) – 用于组合 in_keys 中的子键的分隔符，这些子键是字符串元组。默认为 "_"。

• source (str 或 键列表, 可选) – 如果提供字符串，则指向包含要使用的输入键列表的模块属性。如果改为提供列表，则将包含用作模块输入的键。默认为 "in_keys"，这是 TensorDictModule 的输入键列表的属性名称。

• dest (str 或 键列表, 可选) – 如果提供字符串，则指向包含要使用的输出键列表的模块属性。如果改为提供列表，则将包含用作模块输出的键。默认为 "out_keys"，这是 TensorDictModule 的输出键列表的属性名称。

• auto_batch_size (bool, 可选) – 如果为 True，则输入的 tensordict 的批大小会自动确定为所有输入张量中最大数量的共同维度。默认为 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。通常，建议只对 tensordict 的叶子节点使用 dispatch()。

示例

>>> 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!")
...