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)

tensordict.nn.set_skip_existing¶

文件

教程

資源