tensordict.nn.set_skip_existing

class tensordict.nn.set_skip_existing(mode: bool | None = True, in_key_attr='in_keys', out_key_attr='out_keys')

一个用于跳过 TensorDict 图中现有节点的上下文管理器。

当用作上下文管理器时，它会将 skip_existing() 的值设置为指定的 mode，允许用户编写方法来检查全局值并相应地执行代码。

当用作方法装饰器时，它会检查 tensordict 的输入键，并且如果 skip_existing() 调用返回 True，则在所有输出键都已经存在的情况下跳过该方法。对于不符合以下签名的函数，不应将其用作装饰器：def fun(self, tensordict, *args, **kwargs)。

参数：

• mode (bool, optional) – 如果为 True，表示图中已有的条目不会被覆盖，除非它们只是部分存在。skip_existing() 将返回 True。如果为 False，则不执行任何检查。如果为 None，skip_existing() 的值将不会改变。这旨在专门用于装饰方法，并允许其行为在使用同类作为上下文管理器时取决于上下文管理器（见下例）。默认为 True。

• in_key_attr (str, optional) – 被装饰模块方法中输入键列表属性的名称。默认为 in_keys。

• out_key_attr (str, optional) – 被装饰模块方法中输出键列表属性的名称。默认为 out_keys。

示例

>>> with set_skip_existing():
...     if skip_existing():
...         print("True")
...     else:
...         print("False")
...
True
>>> print("calling from outside:", skip_existing())
calling from outside: False

此类别也可以用作装饰器

示例

>>> from tensordict import TensorDict
>>> from tensordict.nn import set_skip_existing, skip_existing, TensorDictModuleBase
>>> class MyModule(TensorDictModuleBase):
...     in_keys = []
...     out_keys = ["out"]
...     @set_skip_existing()
...     def forward(self, tensordict):
...         print("hello")
...         tensordict.set("out", torch.zeros(()))
...         return tensordict
>>> module = MyModule()
>>> module(TensorDict({"out": torch.zeros(())}, []))  # does not print anything
TensorDict(
    fields={
        out: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> module(TensorDict())  # prints hello
hello
TensorDict(
    fields={
        out: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

将 mode 设置为 None 来装饰方法在希望上下文管理器从外部处理跳过行为时非常有用

示例

>>> from tensordict import TensorDict
>>> from tensordict.nn import set_skip_existing, skip_existing, TensorDictModuleBase
>>> class MyModule(TensorDictModuleBase):
...     in_keys = []
...     out_keys = ["out"]
...     @set_skip_existing(None)
...     def forward(self, tensordict):
...         print("hello")
...         tensordict.set("out", torch.zeros(()))
...         return tensordict
>>> module = MyModule()
>>> _ = module(TensorDict({"out": torch.zeros(())}, []))  # prints "hello"
hello
>>> with set_skip_existing(True):
...     _ = module(TensorDict({"out": torch.zeros(())}, []))  # no print

注意

为了允许模块具有相同的输入和输出键而不错误地忽略子图，每当输出键同时也是输入键时，@set_skip_existing(True) 将被禁用。

>>> class MyModule(TensorDictModuleBase):
...     in_keys = ["out"]
...     out_keys = ["out"]
...     @set_skip_existing()
...     def forward(self, tensordict):
...         print("calling the method!")
...         return tensordict
...
>>> module = MyModule()
>>> module(TensorDict({"out": torch.zeros(())}, []))  # does not print anything
calling the method!
TensorDict(
    fields={
        out: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)