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 实例，则假定 __not__ 正在使用 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!")
...