TensorDictParams¶

class tensordict.TensorDictParams(parameters: TensorDictBase, *, no_convert=False, lock: bool = False)¶

保存一個包含參數的 TensorDictBase 實例。

此類別將包含的參數公開給父 nn.Module，以便迭代模組的參數時，也會迭代 tensordict 的葉節點。

索引的操作方式與封裝的 tensordict 的索引完全相同。參數名稱將使用 flatten_keys("_")() 在此模組中註冊。因此，named_parameters() 的結果和 tensordict 的內容在鍵名方面會略有不同。

在 tensordict 中設定張量的任何操作都會通過 torch.nn.Parameter 轉換來擴充。

參數:

parameters (TensorDictBase) – 表示為參數的 tensordict。除非 no_convert=True，否則值將轉換為參數。

關鍵字參數:

no_convert (bool) – 如果 True，則在建構時和之後都不會發生轉換為 nn.Parameter 的動作（除非 no_convert 屬性已變更）。如果 no_convert 為 True 且存在非參數，則它們將註冊為緩衝區。預設為 False。
lock (bool) – 如果 True，則 TensorDictParams 託管的 tensordict 將會被鎖定。這對於避免不必要的修改很有用，但也限制了可以對物件執行的操作（並且在需要 unlock_() 時可能會產生顯著的效能影響）。預設為 False。

範例

>>> from torch import nn
>>> from tensordict import TensorDict
>>> module = nn.Sequential(nn.Linear(3, 4), nn.Linear(4, 4))
>>> params = TensorDict.from_module(module)
>>> params.lock_()
>>> p = TensorDictParams(params)
>>> print(p)
TensorDictParams(params=TensorDict(
    fields={
        0: TensorDict(
            fields={
                bias: Parameter(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                weight: Parameter(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        1: TensorDict(
            fields={
                bias: Parameter(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                weight: Parameter(shape=torch.Size([4, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False))
>>> class CustomModule(nn.Module):
...     def __init__(self, params):
...         super().__init__()
...         self.params = params
>>> m = CustomModule(p)
>>> # the wrapper supports assignment and values are turned in Parameter
>>> m.params['other'] = torch.randn(3)
>>> assert isinstance(m.params['other'], nn.Parameter)

abs() → T¶: 計算 TensorDict 中每個元素的絕對值。

abs_() → T¶: 就地計算 TensorDict 中每個元素的絕對值。

acos() → T¶: 計算 TensorDict 中每個元素的 acos() 值。

acos_() → T¶: 就地計算 TensorDict 中每個元素的 acos() 值。

add(other: tensordict.base.TensorDictBase | torch.Tensor, *, alpha: Optional[float] = None, default: Optional[Union[str, Tensor]] = None) → TensorDictBase¶

將 other 乘以 alpha 後加到 self。

\[\text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i\]

參數:

other (TensorDictBase 或 torch.Tensor) – 要加到 self 的張量或 TensorDict。

關鍵字參數:

alpha (Number, optional) – other 的乘數。
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

add_(other: tensordict.base.TensorDictBase | float, *, alpha: Optional[float] = None)¶: 原地 (In-place) 版本的 add()。

注意

原地 add 不支援 default 關鍵字引數。

add_module(name: str, module: Optional[Module]) → None¶

將子模組新增到目前模組。

可以使用給定的名稱作為屬性來訪問該模組。

參數:

name (str) – 子模組的名稱。可以使用給定的名稱從此模組訪問子模組
module (Module) – 要新增到模組的子模組。

addcdiv(other1: tensordict.base.TensorDictBase | torch.Tensor, other2: tensordict.base.TensorDictBase | torch.Tensor, value: float | None = 1)¶

執行 other1 除以 other2 的逐元素除法，將結果乘以純量 value，然後將其加到 self。

\[\text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i}\]

self、other1 和 other2 的元素的形狀必須是可廣播的。

對於 FloatTensor 或 DoubleTensor 類型的輸入，value 必須是實數，否則為整數。

參數:

other1 (TensorDict 或 Tensor) – 分子 tensordict（或 tensor）
tensor2 (TensorDict 或 Tensor) – 分母 tensordict（或 tensor）

關鍵字參數:

value (Number, optional) – \(\text{tensor1} / \text{tensor2}\) 的乘數

addcdiv_(other1, other2, *, value: float | None = 1)¶: addcdiv() 的原地 (in-place) 版本。

addcmul(other1, other2, *, value: float | None = 1)¶

執行 other1 乘以 other2 的逐元素乘法，將結果乘以純量 value，然後將其加到 self。

\[\text{out}_i = \text{input}_i + \text{value} \times \text{other1}_i \times \text{other2}_i\]

self、other1 和 other2 的形狀必須是可廣播的。

對於 FloatTensor 或 DoubleTensor 類型的輸入，value 必須是實數，否則為整數。

參數:

other1 (TensorDict 或 Tensor) – 要相乘的 tensordict 或 tensor
other2 (TensorDict 或 Tensor) – 要相乘的 tensordict 或 tensor

關鍵字參數:

value (Number, optional) – \(other1 .* other2\) 的乘數

addcmul_(other1, other2, *, value: float | None = 1)¶: addcmul() 的原地 (in-place) 版本。

all(dim: int = None) → bool | tensordict.base.TensorDictBase¶

檢查 tensordict 中是否所有值都是 True/非空值 (non-null)。

參數:: dim (int, optional) – 若為 None，則回傳一個布林值，指示是否所有張量都回傳 tensor.all() == True。若為整數，則僅當此維度與 tensordict 的形狀相容時，才會對指定的維度呼叫 all。

any(dim: int = None) → bool | tensordict.base.TensorDictBase¶

檢查 tensordict 中是否有任何值為 True/非空值。

參數:: dim (int, optional) – 若為 None，則回傳一個布林值，指示是否所有張量都回傳 tensor.any() == True。若為整數，則僅當此維度與 tensordict 的形狀相容時，才會對指定的維度呼叫 all。

apply(fn: Callable, *others: TensorDictBase, batch_size: Optional[Sequence[int]] = None, device: torch.device | None = _NoDefault.ZERO, names: Optional[Sequence[str]] = _NoDefault.ZERO, inplace: bool = False, default: Any = _NoDefault.ZERO, filter_empty: bool | None = None, call_on_nested: bool = False, **constructor_kwargs) → tensordict.base.TensorDictBase | None¶

將可呼叫物件應用於 tensordict 中儲存的所有值，並將它們設定在新的 tensordict 中。

可呼叫物件的簽章必須為 Callable[Tuple[Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。

參數:

fn (Callable) – 要應用於 tensordict 中的張量的函式。
*others (TensorDictBase 實例, optional) – 如果提供，這些 tensordict 實例應該具有與 self 相同的結構。fn 參數應該接收與 tensordict 數量一樣多的未命名輸入，包括 self。如果其他 tensordict 缺少條目，則可以透過 default 關鍵字參數傳遞預設值。

關鍵字參數:

batch_size (sequence of int, optional) – 如果提供此參數，產生的 TensorDict 將具有期望的 batch_size。 batch_size 參數應與轉換後的 batch_size 相符。這是一個僅限關鍵字的參數。
device (torch.device, optional) – 產生的 device，如果有的話。
names (list of str, optional) – 新的維度名稱，在 batch_size 被修改的情況下。
inplace (bool, optional) – 如果為 True，則會進行原地修改。預設值為 False。這是一個僅限關鍵字的參數。
default (Any, optional) – 其他 tensordict 中缺失條目的預設值。如果未提供，缺失的條目將引發 KeyError。
filter_empty (bool, optional) – 如果 True，空的 tensordict 將會被過濾掉。這也降低了計算成本，因為不會建立和銷毀空的資料結構。非張量資料會被視為葉節點，因此即使該函數沒有更動到，仍會保留在 tensordict 中。為了向後相容性，預設值為 False。
propagate_lock (bool, optional) – 如果 True，被鎖定的 tensordict 會產生另一個被鎖定的 tensordict。預設值為 False。

call_on_nested (bool, optional) –

如果 True，該函數將在第一層張量和容器（TensorDict 或 tensorclass）上呼叫。在這種情況下，func 負責將其呼叫傳播到巢狀層級。這允許在將呼叫傳播到巢狀 tensordict 時進行細粒度的行為控制。如果 False，該函數只會在葉節點上呼叫，並且 apply 會負責將該函數分派到所有葉節點。

>>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]})
>>> def mean_tensor_only(val):
...     if is_tensor_collection(val):
...         raise RuntimeError("Unexpected!")
...     return val.mean()
>>> td_mean = td.apply(mean_tensor_only)
>>> def mean_any(val):
...     if is_tensor_collection(val):
...         # Recurse
...         return val.apply(mean_any, call_on_nested=True)
...     return val.mean()
>>> td_mean = td.apply(mean_any, call_on_nested=True)

out (TensorDictBase, optional) –
一個用於寫入結果的 tensordict。這可以用來避免創建新的 tensordict。
```
>>> td = TensorDict({"a": 0})
>>> td.apply(lambda x: x+1, out=td)
>>> assert (td==1).all()
```
警告

如果在 tensordict 上執行的操作需要訪問多個鍵才能進行單一計算，則提供一個等於 self 的 out 參數可能會導致該操作默默地提供錯誤的結果。例如
```
>>> td = TensorDict({"a": 1, "b": 1})
>>> td.apply(lambda x: x+td["a"])["b"] # Right!
tensor(2)
>>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong!
tensor(3)
```
**constructor_kwargs – 要傳遞給 TensorDict 建構子的其他關鍵字參數。

回傳值:

一個帶有 transformed_in 張量的新 tensordict。

範例

>>> td = TensorDict({
...     "a": -torch.ones(3),
...     "b": {"c": torch.ones(3)}},
...     batch_size=[3])
>>> td_1 = td.apply(lambda x: x+1)
>>> assert (td_1["a"] == 0).all()
>>> assert (td_1["b", "c"] == 2).all()
>>> td_2 = td.apply(lambda x, y: x+y, td)
>>> assert (td_2["a"] == -2).all()
>>> assert (td_2["b", "c"] == 2).all()

注意

如果函數回傳 None，則會忽略該條目。這可以用來過濾 tensordict 中的資料。

>>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, [])
>>> def filter(tensor):
...     if tensor == 1:
...         return tensor
>>> td.apply(filter)
TensorDict(
    fields={
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

注意

apply 方法將回傳一個 TensorDict 實例，無論輸入類型為何。要保持相同的類型，可以執行

>>> out = td.clone(False).update(td.apply(...))

apply_(fn: Callable, *others, **kwargs) → T¶

將一個可呼叫物件應用於 tensordict 中儲存的所有值，並將它們原地重寫。

參數:

fn (Callable) – 要應用於 tensordict 中的張量的函式。
*others (sequence of TensorDictBase, optional) – 要使用的其他 tensordict。

關鍵字參數：請參閱 apply()。

回傳值:: 套用該函數的 self 或 self 的副本

asin() → T¶: 計算 TensorDict 中每個元素的 asin() 值。

asin_() → T¶: 原地計算 TensorDict 中每個元素的 asin() 值。

atan() → T¶: 計算 TensorDict 中每個元素的 atan() 值。

atan_() → T¶: 原地計算 TensorDict 中每個元素的 atan() 值。

auto_batch_size_(batch_dims: Optional[int] = None) → T¶

設定 tensordict 的最大 batch-size，最多到可選的 batch_dims。

參數:: batch_dims (int, optional) – 如果提供，batch-size 最多為 batch_dims 的長度。
回傳值:: self

範例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict({"a": torch.randn(3, 4, 5), "b": {"c": torch.randn(3, 4, 6)}}, batch_size=[])
>>> td.auto_batch_size_()
>>> print(td.batch_size)
torch.Size([3, 4])
>>> td.auto_batch_size_(batch_dims=1)
>>> print(td.batch_size)
torch.Size([3])

auto_device_() → T¶

自動設定裝置 (device)，如果它是唯一的。

傳回：具有已編輯的 device 屬性的 self。

property batch_dims: int¶

tensordict batch size 的長度。

回傳值:: 描述 tensordict 維度的數量的整數。

property batch_size: Size¶

TensorDict 的形狀 (shape) (或 batch_size)。

tensordict 的形狀對應於其包含的 tensors 的共同前 N 維度，其中 N 是一個任意數字。batch-size 與代表 tensor 語義相關形狀的“特徵大小 (feature size)”形成對比。例如，一批影片可能具有 [B, T, C, W, H] 的形狀，其中 [B, T] 是 batch-size（batch 和時間維度），而 [C, W, H] 是特徵維度（通道和空間維度）。

TensorDict 形狀由使用者在初始化時控制（即，它不是從 tensor 形狀推斷出來的）。

如果新大小與 TensorDict 內容相容，則可以動態編輯 batch_size。例如，始終允許將 batch size 設定為空值。

回傳值:: 描述 TensorDict batch size 的 Size 物件。

範例

>>> data = TensorDict({
...     "key 0": torch.randn(3, 4),
...     "key 1": torch.randn(3, 5),
...     "nested": TensorDict({"key 0": torch.randn(3, 4)}, batch_size=[3, 4])},
...     batch_size=[3])
>>> data.batch_size = () # resets the batch-size to an empty value

bfloat16()¶: 將所有 tensors 轉換為 torch.bfloat16。

bool()¶: 將所有 tensors 轉換為 torch.bool。

buffers(recurse: bool = True) → Iterator[Tensor]¶

傳回 module buffers 的迭代器。

參數:: recurse (bool) – 如果為 True，則產生此 module 和所有子 modules 的 buffers。否則，僅產生作為此 module 直接成員的 buffers。
產生：: torch.Tensor – module buffer

範例

>>> # xdoctest: +SKIP("undefined vars")
>>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

bytes(*, count_duplicates: bool = True) → int¶

計算所包含 tensors 的位元組數。

關鍵字參數:: count_duplicates (bool) – 是否將重複的 tensor 計算為獨立的 tensor。如果 False，則僅會丟棄嚴格相同的 tensors（來自共同基礎 tensor 的相同檢視但具有不同的 ids 將被計算兩次）。預設為 True（假設每個 tensor 都是單一副本）。

classmethod cat(input, dim=0, *, out=None)¶

沿著給定的維度將 tensordict 連接成一個單一的 tensordict。

此呼叫等同於呼叫 torch.cat()，但與 torch.compile 相容。

cat_from_tensordict(dim: int = 0, *, sorted: Optional[Union[bool, List[NestedKey]]] = None, out: Tensor = None) → Tensor¶

將 tensordict 的所有條目連接成一個單一的張量。

參數:

dim (int, optional) – 條目應沿著此維度連接。

關鍵字參數:

sorted (bool 或 NestedKeys 的列表) – 如果 True，條目將按字母順序連接。如果 False（預設），將使用 dict 順序。或者，可以提供金鑰名稱的列表，並且張量將相應地連接。這會產生一些額外的負擔，因為金鑰列表將針對 tensordict 中的葉節點名稱列表進行檢查。
out (torch.Tensor, optional) – 用於 cat 運算的選擇性目標張量。

cat_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → T¶

將條目連接到一個新條目中，並可能移除原始值。

參數:: keys (NestedKey 的序列) – 要連接的條目。

關鍵字參數

out_key (NestedKey): 連接輸入的新金鑰名稱。 keep_entries (bool, optional): 如果 False, keys 中的條目將被刪除。

預設為 False。

dim (int, optional): 必須進行連接的維度。: 預設為 0。

Returns: self

範例

>>> td = TensorDict(a=torch.zeros(1), b=torch.ones(1))
>>> td.cat_tensors("a", "b", out_key="c")
>>> assert "a" not in td
>>> assert (td["c"] == torch.tensor([0, 1])).all()

ceil() → T¶: 計算 TensorDict 中每個元素的 ceil() 值。

ceil_() → T¶: 就地計算 TensorDict 中每個元素的 ceil() 值。

children() → Iterator[Module]¶

傳回一個迭代器，用於迭代直接子模組。

產生：: Module – 一個子模組

chunk(chunks: int, dim: int = 0) → tuple[tensordict.base.TensorDictBase, ...]¶

若可能，將 tensordict 分割成指定數量的區塊。

每個區塊都是輸入 tensordict 的一個視圖。

參數:

chunks (int) – 要傳回的區塊數量
dim (int, optional) – 要沿其分割 tensordict 的維度。預設值為 0。

範例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> td0, td1 = td.chunk(dim=-1, chunks=2)
>>> td0['x']
tensor([[[ 0,  1],
         [ 2,  3]],
        [[ 8,  9],
         [10, 11]],
        [[16, 17],
         [18, 19]]])

clamp_max(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T¶

如果 self 的元素大於 other 的值，則將其鉗制（clamp）到 other。

參數:: other (TensorDict or Tensor) – 另一個輸入 tensordict 或 tensor。
關鍵字參數:: default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

clamp_max_(other: tensordict.base.TensorDictBase | torch.Tensor) → T¶: clamp_max() 的原地（in-place）版本。

注意

原地 clamp_max 不支援 default 關鍵字引數。

clamp_min(other: tensordict.base.TensorDictBase | torch.Tensor, default: Optional[Union[str, Tensor]] = None) → T¶

如果 self 的元素小於 other 的值，則將其鉗制（clamp）到 other。

參數:: other (TensorDict or Tensor) – 另一個輸入 tensordict 或 tensor。
關鍵字參數:: default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

clamp_min_(other: tensordict.base.TensorDictBase | torch.Tensor) → T¶: 原地 (In-place) 版本的 clamp_min()。

注意

原地 (Inplace) clamp_min 不支援 default 關鍵字引數。

clear() → T¶: 清除 tensordict 的內容。

clear_device_() → T¶

清除 tensordict 的裝置 (device)。

Returns: self

clone(recurse: bool = True, **kwargs) → T¶

將 TensorDictBase 子類別的實例複製 (clone) 到相同類型的新 TensorDictBase 子類別中。

若要從任何其他 TensorDictBase 子類型建立 TensorDict 實例，請改為呼叫 to_tensordict() 方法。

參數:: recurse (bool, optional) – 如果 True，TensorDict 中包含的每個張量 (tensor) 也會被複製。否則，只會複製 TensorDict 的樹狀結構。預設值為 True。

注意

與許多其他運算 (逐點運算、形狀運算等) 不同，clone 不會繼承原始的 lock 屬性。這種設計選擇是為了可以創建一個克隆來進行修改，這是最常見的用法。

compile(*args, **kwargs)¶

使用 torch.compile() 編譯此 Module 的 forward。

此 Module 的 __call__ 方法被編譯，並且所有引數都按原樣傳遞給 torch.compile()。

有關此函數引數的詳細資訊，請參閱 torch.compile()。

complex128()¶: 將所有張量轉換為 torch.complex128。

complex32()¶: 將所有張量轉換為 torch.complex32。

complex64()¶: 將所有張量轉換為 torch.complex64。

consolidate(filename: Optional[Union[Path, str]] = None, *, num_threads=0, device: Optional[device] = None, non_blocking: bool = False, inplace: bool = False, return_early: bool = False, use_buffer: bool = False, share_memory: bool = False, pin_memory: bool = False, metadata: bool = False) → None¶

將 tensordict 的內容整合到單一儲存空間中，以實現快速序列化。

參數:

filename (Path, optional) – 一個可選的檔案路徑，用於作為 tensordict 的儲存空間的記憶體對應張量。

關鍵字參數:

num_threads (integer, optional) – 用於填充儲存空間的執行緒數量。
device (torch.device, optional) – 一個可選的裝置，儲存空間必須在此裝置上實例化。
non_blocking (bool, optional) – 傳遞給 copy_() 的 non_blocking 參數。
inplace (bool, optional) – 如果 True，則產生的 tensordict 與更新值的 self 相同。預設為 False。
return_early (bool, optional) – 如果 True 且 num_threads>0，則該方法將傳回 tensordict 的 future。可以使用 future.result() 查詢產生的 tensordict。
use_buffer (bool, optional) – 如果 True 且傳遞了檔名，則會在共享記憶體中建立一個中間本地緩衝區，並且資料將作為最後一個步驟複製到儲存位置。這可能比直接寫入遠端物理記憶體（例如，NFS）更快。預設為 False。
share_memory (bool, optional) – 如果 True，則儲存空間將放置在共享記憶體中。預設為 False。
pin_memory (bool, optional) – 是否應將整合的資料放置在釘選記憶體中。預設為 False。
metadata (bool, optional) – 如果 True，則中繼資料將與通用儲存空間一起儲存。如果提供了檔名，則此操作無效。當人們想要控制如何實現序列化時，儲存中繼資料可能很有用，因為如果中繼資料可用或不可用，TensorDict 會以不同的方式處理整合 TD 的 pickling/unpickling。

注意

如果 tensordict 已經整合，則所有參數都會被忽略，並傳回 self。呼叫 contiguous() 以重新整合。

範例

>>> import pickle
>>> import tempfile
>>> import torch
>>> import tqdm
>>> from torch.utils.benchmark import Timer
>>> from tensordict import TensorDict
>>> data = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> data_consolidated = data.consolidate()
>>> # check that the data has a single data_ptr()
>>> assert torch.tensor([
...     v.untyped_storage().data_ptr() for v in data_c.values(True, True)
... ]).unique().numel() == 1
>>> # Serializing the tensordict will be faster with data_consolidated
>>> with open("data.pickle", "wb") as f:
...    print("regular", Timer("pickle.dump(data, f)", globals=globals()).adaptive_autorange())
>>> with open("data_c.pickle", "wb") as f:
...     print("consolidated", Timer("pickle.dump(data_consolidated, f)", globals=globals()).adaptive_autorange())

contiguous(*args, **kwargs)¶: 傳回一個具有連續值（如果值已經連續，則傳回 self）的相同類型的新 tensordict。

copy()¶

傳回 tensordict 的淺層副本（即，複製結構但不複製資料）。

等效於 TensorDictBase.clone(recurse=False)

copy_(tensordict: T, non_blocking: bool = None) → T¶

請參閱 TensorDictBase.update_。

non-blocking 引數將被忽略，它的存在僅是為了與 torch.Tensor.copy_() 相容。

copy_at_(tensordict: T, idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], non_blocking: bool = False) → T¶: 請參閱 TensorDictBase.update_at_。

cos() → T¶: 計算 TensorDict 中每個元素的 cos() 值。

cos_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 cos() 值。

cosh() → T¶: 計算 TensorDict 中每個元素的 cosh() 值。

cosh_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 cosh() 值。

cpu()¶

將 tensordict 轉換到 CPU。

此函數也支援 to() 的所有關鍵字引數。

create_nested(key)¶

建立一個具有與目前 tensordict 相同形狀、裝置和維度名稱的巢狀 tensordict。

如果該值已存在，則會被此操作覆寫。此操作在鎖定的 tensordict 中會被阻擋。

範例

>>> data = TensorDict({}, [3, 4, 5])
>>> data.create_nested("root")
>>> data.create_nested(("some", "nested", "value"))
>>> print(data)
TensorDict(
    fields={
        root: TensorDict(
            fields={
            },
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False),
        some: TensorDict(
            fields={
                nested: TensorDict(
                    fields={
                        value: TensorDict(
                            fields={
                            },
                            batch_size=torch.Size([3, 4, 5]),
                            device=None,
                            is_shared=False)},
                    batch_size=torch.Size([3, 4, 5]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4, 5]),
    device=None,
    is_shared=False)

cuda(device=None)¶

將 tensordict 轉換到 cuda 裝置（如果尚未在其上）。

參數:: device (int, optional) – 如果提供，則指定 tensor 應轉換到的 cuda 裝置。

此函數也支援 to() 的所有關鍵字引數。

property data¶: 傳回一個包含葉 tensors 的 .data 屬性的 tensordict。

data_ptr(*, storage: bool = False)¶

傳回 tensordict 葉節點的 data_ptr。

這可用於檢查兩個 tensordict 是否共享相同的 data_ptr()。

關鍵字參數:: storage (bool, optional) – 如果為 True，則會呼叫 tensor.untyped_storage().data_ptr()。預設為 False。

範例

>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> assert (td0.data_ptr() == td.data_ptr()).all()

注意

LazyStackedTensorDict 實例會顯示為巢狀 tensordict，以反映其葉節點的真實 data_ptr()

>>> td0 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> td1 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> td = TensorDict.lazy_stack([td0, td1])
>>> td.data_ptr()
TensorDict(
    fields={
        0: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False),
        1: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

del_(*args, **kwargs)¶

刪除 tensordict 的鍵。

參數:: key (NestedKey) – 要刪除的鍵
回傳值:: self

densify(layout: layout = torch.strided)¶

嘗試使用連續張量（普通張量或巢狀張量）表示 lazy stack。

關鍵字參數:: layout (torch.layout) – 巢狀張量的 layout（如果有的話）。預設為 strided。

property depth: int¶

傳回 tensordict 的深度（最大層數）。

最小深度為 0（沒有巢狀 tensordict）。

detach() → T¶

分離 tensordict 中的張量。

回傳值:: 一個新的 tensordict，沒有需要梯度的張量。

detach_(*args, **kwargs)¶

就地分離 tensordict 中的張量。

回傳值:: self。

property device¶

TensorDict 的裝置。

如果 TensorDict 有指定的裝置，則其所有張量（包括巢狀張量）都必須位於同一裝置上。如果 TensorDict 裝置為 None，則不同的值可以位於不同的裝置上。

回傳值:: torch.device 物件，指示張量所在的裝置，如果 TensorDict 沒有裝置，則為 None。

範例

>>> td = TensorDict({
...     "cpu": torch.randn(3, device='cpu'),
...     "cuda": torch.randn(3, device='cuda'),
... }, batch_size=[], device=None)
>>> td['cpu'].device
device(type='cpu')
>>> td['cuda'].device
device(type='cuda')
>>> td = TensorDict({
...     "x": torch.randn(3, device='cpu'),
...     "y": torch.randn(3, device='cuda'),
... }, batch_size=[], device='cuda')
>>> td['x'].device
device(type='cuda')
>>> td['y'].device
device(type='cuda')
>>> td = TensorDict({
...     "x": torch.randn(3, device='cpu'),
...     "y": TensorDict({'z': torch.randn(3, device='cpu')}, batch_size=[], device=None),
... }, batch_size=[], device='cuda')
>>> td['x'].device
device(type='cuda')
>>> td['y'].device # nested tensordicts are also mapped onto the appropriate device.
device(type='cuda')
>>> td['y', 'x'].device
device(type='cuda')

dim() → int¶: 請參閱 batch_dims()。

div(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T¶

將輸入 self 的每個元素除以 other 的對應元素。

\[\text{out}_i = \frac{\text{input}_i}{\text{other}_i}\]

支援廣播、類型提升和整數、浮點數、tensordict 或張量輸入。始終將整數類型提升為預設純量類型。

參數:: other (TensorDict, Tensor 或 Number) – 除數。
關鍵字參數:: default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

div_(other: tensordict.base.TensorDictBase | torch.Tensor) → T¶: div() 的就地版本。

注意

就地 div 不支援 default 關鍵字引數。

double()¶: 將所有 tensors 轉換為 torch.bool。

property dtype¶: 傳回 tensordict 中值的 dtype（如果它是唯一的）。

dumps(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

將 tensordict 儲存到磁碟。

此函式是 memmap() 的代理。

empty(recurse=False, *, batch_size=None, device=_NoDefault.ZERO, names=None) → T¶

傳回具有相同裝置和批次大小的新且空的 tensordict。

參數:

recurse (bool, optional) – 如果 True，則會複製 TensorDict 的完整結構而不包含內容。否則，只會複製根目錄。預設值為 False。

關鍵字參數:

batch_size (torch.Size, optional) – tensordict 的新批次大小。
device (torch.device, optional) – 新裝置。
names (list of str, optional) – 維度名稱。

entry_class(*args, **kwargs)¶

傳回 entry 的類別，可能會避免呼叫 isinstance(td.get(key), type)。

只要 get() 的執行成本很高，就應該優先使用此方法，而不是 tensordict.get(key).shape。

erf() → T¶: 計算 TensorDict 中每個元素的 erf() 值。

erf_() → T¶: 就地計算 TensorDict 中每個元素的 erf() 值。

erfc() → T¶: 計算 TensorDict 中每個元素的 erfc() 值。

erfc_() → T¶: 就地計算 TensorDict 中每個元素的 erfc() 值。

eval() → T¶

將模組設定為評估模式。

這僅對某些模組有效。請參閱特定模組的文件，了解它們在訓練/評估模式下的行為詳細資訊（如果它們受到影響），例如 Dropout、BatchNorm 等。

這等同於 self.train(False)。

請參閱在本地停用梯度計算，以比較 .eval() 與幾個可能與其混淆的類似機制。

回傳值:: self
回傳類型:: 模組

exclude(*keys: NestedKey, inplace: bool = False) → T¶

排除 tensordict 中的鍵，並返回一個不包含這些條目的新 tensordict。

數值不會被複製：對原始或新的 tensordict 的張量進行原地 (in-place) 修改將會導致兩個 tensordict 都發生變化。

參數:

*keys (str) – 要排除的鍵。
inplace (bool) – 如果為 True，則 tensordict 會原地修剪。預設值為 False。

回傳值:

一個新的 tensordict (如果 inplace=True 則為同一個) ，不包含排除的條目。

範例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, [])
>>> td.exclude("a", ("b", "c"))
TensorDict(
    fields={
        b: TensorDict(
            fields={
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.exclude("a", "b")
TensorDict(
    fields={
    },
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

exp() → T¶: 計算 TensorDict 中每個元素的 exp() 值。

exp_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 exp() 值。

expand(*args, **kwargs) → T¶

根據 expand() 函數展開 tensordict 的每個張量，忽略特徵維度。

支援使用可迭代物件來指定形狀。

範例

>>> td = TensorDict({
...     'a': torch.zeros(3, 4, 5),
...     'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4])
>>> td_expand = td.expand(10, 3, 4)
>>> assert td_expand.shape == torch.Size([10, 3, 4])
>>> assert td_expand.get("a").shape == torch.Size([10, 3, 4, 5])

expand_as(other: tensordict.base.TensorDictBase | torch.Tensor) → TensorDictBase¶

將 tensordict 的形狀廣播到 other 的形狀，並相應地展開它。

如果輸入是張量集合（tensordict 或 tensorclass），則葉節點將以一對一的方式展開。

範例

>>> from tensordict import TensorDict
>>> import torch
>>> td0 = TensorDict({
...     "a": torch.ones(3, 1, 4),
...     "b": {"c": torch.ones(3, 2, 1, 4)}},
...     batch_size=[3],
... )
>>> td1 = TensorDict({
...     "a": torch.zeros(2, 3, 5, 4),
...     "b": {"c": torch.zeros(2, 3, 2, 6, 4)}},
...     batch_size=[2, 3],
... )
>>> expanded = td0.expand_as(td1)
>>> assert (expanded==1).all()
>>> print(expanded)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([2, 3, 5, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([2, 3, 2, 6, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([2, 3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([2, 3]),
    device=None,
    is_shared=False)

expm1() → T¶: 計算 TensorDict 中每個元素的 expm1() 值。

expm1_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 expm1() 值。

extra_repr() → str¶

設定模組的額外表示。

若要列印客製化的額外資訊，您應該在自己的模組中重新實作此方法。單行和多行字串都是可接受的。

fill_(key: NestedKey, value: float | bool) → T¶

使用給定的純量值填滿鍵所指向的張量。

參數:

key (str 或 巢狀鍵) – 要填滿的條目。
value (Number 或 bool) – 用於填滿的值。

回傳值:

self

filter_empty_()¶: 原地 (in-place) 過濾掉所有空的 tensordict。

filter_non_tensor_data() → T¶: 過濾掉所有非張量資料 (non-tensor-data)。

flatten(start_dim=0, end_dim=- 1)¶

展平 (flatten) tensordict 的所有張量。

參數:

start_dim (int) – 要展平的第一個維度
end_dim (int) – 要展平的最後一個維度

範例

>>> td = TensorDict({
...     "a": torch.arange(60).view(3, 4, 5),
...     "b": torch.arange(12).view(3, 4)}, batch_size=[3, 4])
>>> td_flat = td.flatten(0, 1)
>>> td_flat.batch_size
torch.Size([12])
>>> td_flat["a"]
tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44],
        [45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59]])
>>> td_flat["b"]
tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])

flatten_keys(separator: str = '.', inplace: bool = False) → TensorDictBase¶

將巢狀的 tensordict 遞迴地轉換為扁平的 tensordict。

TensorDict 的類型將會遺失，結果會是一個簡單的 TensorDict 實例。

參數:

separator (str, optional) – 巢狀項目之間的分隔符。
inplace (bool, optional) – 如果 True，則產生的 tensordict 將與呼叫它的 tensordict 具有相同的 identity。預設為 False。
is_leaf (callable, optional) – 一個可呼叫物件，作用於類別類型，並傳回一個布林值，指示此類別是否應被視為葉節點。

範例

>>> data = TensorDict({"a": 1, ("b", "c"): 2, ("e", "f", "g"): 3}, batch_size=[])
>>> data.flatten_keys(separator=" - ")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b - c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        e - f - g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

此方法和 unflatten_keys() 在處理 state-dicts 時特別有用，因為它們可以無縫地將扁平的字典轉換為模仿模型結構的資料結構。

範例

>>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4))
>>> ddp_model = torch.ao.quantization.QuantWrapper(model)
>>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".")
>>> print(state_dict)
TensorDict(
    fields={
        module: TensorDict(
            fields={
                0: TensorDict(
                    fields={
                        bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                        weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> model_state_dict = state_dict.get("module")
>>> print(model_state_dict)
TensorDict(
    fields={
        0: TensorDict(
            fields={
                bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> model.load_state_dict(dict(model_state_dict.flatten_keys(".")))

float()¶: 將所有 tensors 轉換為 torch.float。

float16()¶: 將所有 tensors 轉換為 torch.float16。

float32()¶: 將所有 tensors 轉換為 torch.float32。

float64()¶: 將所有 tensors 轉換為 torch.float64。

floor() → T¶: 計算 TensorDict 中每個元素的 floor() 值。

floor_() → T¶: 就地計算 TensorDict 中每個元素的 floor() 值。

forward(*input: Any) → None¶

定義每次呼叫時執行的計算。

應被所有子類別覆寫。

注意

雖然 forward pass 的配方需要在這個函式中定義，但應該在之後呼叫 Module 實例，而不是這個函式，因為前者會處理已註冊的 hooks，而後者會靜默地忽略它們。

frac() → T¶: 計算 TensorDict 中每個元素的 frac() 值。

frac_() → T¶: 就地計算 TensorDict 中每個元素的 frac() 值。

classmethod from_dict(*args, **kwargs)¶

從字典或另一個 TensorDict 建立並傳回一個 TensorDict。

如果未指定 batch_size，則傳回可能的最大 batch size。

此函式也適用於巢狀字典，或可用於確定巢狀 tensordict 的 batch-size。

參數:

input_dict (dictionary, optional) – 用作資料來源的字典（相容巢狀鍵）。
batch_size (iterable of int, optional) – tensordict 的 batch size。
device (torch.device or compatible type, optional) – TensorDict 的裝置。
batch_dims (int, optional) – batch_dims (即要考量為 batch_size 的前導維度的數量)。與 batch_size 互斥。請注意，這是 tensordict 的 __maximum__ 批次維度數量，容許較小的數量。
names (list of str, optional) – tensordict 的維度名稱列表。

範例

>>> input_dict = {"a": torch.randn(3, 4), "b": torch.randn(3)}
>>> print(TensorDict.from_dict(input_dict))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> # nested dict: the nested TensorDict can have a different batch-size
>>> # as long as its leading dims match.
>>> input_dict = {"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}
>>> print(TensorDict.from_dict(input_dict))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> # we can also use this to work out the batch sie of a tensordict
>>> input_td = TensorDict({"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}, [])
>>> print(TensorDict.from_dict(input_td))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

from_dict_instance(input_dict, batch_size=None, device=None, batch_dims=None)¶

from_dict() 的實例方法版本。

與 from_dict() 不同，此方法會嘗試將 tensordict 類型保留在現有樹狀結構中 (對於任何現有的葉節點)。

範例

>>> from tensordict import TensorDict, tensorclass
>>> import torch
>>>
>>> @tensorclass
>>> class MyClass:
...     x: torch.Tensor
...     y: int
>>>
>>> td = TensorDict({"a": torch.randn(()), "b": MyClass(x=torch.zeros(()), y=1)})
>>> print(td.from_dict_instance(td.to_dict()))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: MyClass(
            x=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
            y=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(td.from_dict(td.to_dict()))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                y: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

classmethod from_h5(filename, mode='r')¶

從 h5 檔案建立 PersistentTensorDict。

此函數將自動決定每個巢狀 tensordict 的批次大小。

參數:

filename (str) – h5 檔案的路徑。
mode (str, optional) – 讀取模式。預設為 "r"。

classmethod from_module(module, as_module: bool = False, lock: bool = True, use_state_dict: bool = False)¶

將模組的參數和緩衝區複製到 tensordict 中。

參數:

module (nn.Module) – 要從中取得參數的模組。
as_module (bool, optional) – 如果為 True，將返回 TensorDictParams 實例，該實例可用於將參數儲存在 torch.nn.Module 中。預設為 False。
lock (bool, optional) – 如果為 True，則結果 tensordict 將被鎖定。預設為 True。
use_state_dict (bool, optional) –
如果為 True，將使用模組的 state-dict 並將其解壓成具有模型樹狀結構的 TensorDict。預設為 False。

注意

當必須使用 state-dict hooks 時，這特別有用。

範例

>>> from torch import nn
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1
... )
>>> params = TensorDict.from_module(module)
>>> print(params["layers", "0", "linear1"])
TensorDict(
    fields={
        bias: Parameter(shape=torch.Size([2048]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Parameter(shape=torch.Size([2048, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

classmethod from_modules(*modules, as_module: bool = False, lock: bool = True, use_state_dict: bool = False, lazy_stack: bool = False, expand_identical: bool = False)¶

透過 vmap 檢索多個模組的參數，以進行集成學習/期望應用程式的功能。

參數:

modules (nn.Module 的序列) – 要從中取得參數的模組序列。如果模組的結構不同，則需要使用延遲堆疊（請參閱下面的 lazy_stack 參數）。

關鍵字參數:

as_module (bool, 可選) – 如果為 True，將會回傳一個 TensorDictParams 實例，可用於在 torch.nn.Module 中儲存參數。預設為 False。
lock (bool, 可選) – 如果為 True，則產生的 tensordict 將會被鎖定。預設為 True。
use_state_dict (bool, 可選) –
如果為 True，將使用模組的 state-dict 並將其解壓成具有模型樹狀結構的 TensorDict。預設為 False。

注意

當必須使用 state-dict hooks 時，這特別有用。
lazy_stack (bool, 可選) –
參數應該被密集堆疊還是延遲堆疊。預設為 False (密集堆疊)。

注意

lazy_stack 和 as_module 是互斥的功能。

警告

延遲輸出和非延遲輸出之間有一個關鍵差異，即非延遲輸出將使用所需的批次大小重新實例化參數，而 lazy_stack 只會將參數表示為延遲堆疊。這表示當 lazy_stack=True 時，原始參數可以安全地傳遞給優化器，但當它設定為 True 時，需要傳遞新的參數。

警告

雖然使用延遲堆疊來保持原始參數參考可能很誘人，但請記住，延遲堆疊會在每次呼叫 get() 時執行堆疊。這將需要記憶體（參數大小的 N 倍，如果建立圖形則更多）和計算時間。這也意味著優化器將包含更多參數，並且諸如 step() 或 zero_grad() 之類的操作將需要更長的時間才能執行。通常，lazy_stack 應該僅保留給極少數的使用情況。
expand_identical (bool, 可選) – 如果為 True 並且相同的參數（相同的識別碼）正在堆疊到自身，則將會回傳該參數的擴展版本。當 lazy_stack=True 時，會忽略此參數。

範例

>>> from torch import nn
>>> from tensordict import TensorDict
>>> torch.manual_seed(0)
>>> empty_module = nn.Linear(3, 4, device="meta")
>>> n_models = 2
>>> modules = [nn.Linear(3, 4) for _ in range(n_models)]
>>> params = TensorDict.from_modules(*modules)
>>> print(params)
TensorDict(
    fields={
        bias: Parameter(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Parameter(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)
>>> # example of batch execution
>>> def exec_module(params, x):
...     with params.to_module(empty_module):
...         return empty_module(x)
>>> x = torch.randn(3)
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> assert y.shape == (n_models, 4)
>>> # since lazy_stack = False, backprop leaves the original params untouched
>>> y.sum().backward()
>>> assert params["weight"].grad.norm() > 0
>>> assert modules[0].weight.grad is None

使用 lazy_stack=True，情況略有不同

>>> params = TensorDict.from_modules(*modules, lazy_stack=True)
>>> print(params)
LazyStackedTensorDict(
    fields={
        bias: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Tensor(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    exclusive_fields={
    },
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False,
    stack_dim=0)
>>> # example of batch execution
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> assert y.shape == (n_models, 4)
>>> y.sum().backward()
>>> assert modules[0].weight.grad is not None

classmethod from_namedtuple(named_tuple, *, auto_batch_size: bool = False)¶

將 namedtuple 遞迴地轉換為 TensorDict。

關鍵字參數:: auto_batch_size (bool, 可選) – 如果為 True，將會自動計算批次大小。預設為 False。

範例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> nt = data.to_namedtuple()
>>> print(nt)
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))
>>> TensorDict.from_namedtuple(nt, auto_batch_size=True)
TensorDict(
    fields={
        a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                a_string: NonTensorData(data=zero!, batch_size=torch.Size([3]), device=None),
                a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

classmethod from_pytree(pytree, *, batch_size: Optional[Size] = None, auto_batch_size: bool = False, batch_dims: Optional[int] = None)¶

將 pytree 轉換為 TensorDict 實例。

此方法旨在盡可能保持 pytree 的巢狀結構。

會新增額外的非張量鍵，以追蹤每個層級的識別碼，從而提供內建的 pytree-to-tensordict 雙射轉換 API。

目前接受的類別包括 lists、tuples、named tuples 和 dict。

注意

對於字典，非 NestedKey 鍵會作為 NonTensorData 實例單獨註冊。

注意

可轉換為張量的類型 (例如 int、float 或 np.ndarray) 將會轉換為 torch.Tensor 實例。請注意，此轉換是滿射的：將 tensordict 轉換回 pytree 不會恢復原始類型。

範例

>>> # Create a pytree with tensor leaves, and one "weird"-looking dict key
>>> class WeirdLookingClass:
...     pass
...
>>> weird_key = WeirdLookingClass()
>>> # Make a pytree with tuple, lists, dict and namedtuple
>>> pytree = (
...     [torch.randint(10, (3,)), torch.zeros(2)],
...     {
...         "tensor": torch.randn(
...             2,
...         ),
...         "td": TensorDict({"one": 1}),
...         weird_key: torch.randint(10, (2,)),
...         "list": [1, 2, 3],
...     },
...     {"named_tuple": TensorDict({"two": torch.ones(1) * 2}).to_namedtuple()},
... )
>>> # Build a TensorDict from that pytree
>>> td = TensorDict.from_pytree(pytree)
>>> # Recover the pytree
>>> pytree_recon = td.to_pytree()
>>> # Check that the leaves match
>>> def check(v1, v2):
>>>     assert (v1 == v2).all()
>>>
>>> torch.utils._pytree.tree_map(check, pytree, pytree_recon)
>>> assert weird_key in pytree_recon[1]

classmethod from_struct_array(struct_array: ndarray, device: Optional[device] = None) → T¶

將結構化的 numpy 陣列轉換為 TensorDict。

產生的 TensorDict 的內容將與 numpy 陣列共享相同的記憶體內容（這是一個零複製操作）。就地更改結構化的 numpy 陣列的值將會影響 TensorDict 的內容。

範例

>>> x = np.array(
...     [("Rex", 9, 81.0), ("Fido", 3, 27.0)],
...     dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")],
... )
>>> td = TensorDict.from_struct_array(x)
>>> x_recon = td.to_struct_array()
>>> assert (x_recon == x).all()
>>> assert x_recon.shape == x.shape
>>> # Try modifying x age field and check effect on td
>>> x["age"] += 1
>>> assert (td["age"] == np.array([10, 4])).all()

classmethod fromkeys(keys: List[NestedKey], value: Any = 0)¶

從鍵列表和單個值建立一個 tensordict。

參數:

keys (NestedKey 的列表) – 一個可迭代物件，指定新字典的鍵。
value (相容類型, 選用) – 所有鍵的值。預設為 0。

gather(dim: int, index: Tensor, out: Optional[T] = None) → T¶

沿著由 dim 指定的軸收集值。

參數:

dim (int) – 收集元素沿著的維度
index (torch.Tensor) – 一個長整數張量，其維度數量與 tensordict 的維度數量匹配，並且兩者之間只有一個維度不同（即收集維度）。它的元素指向要沿所需維度收集的索引。
out (TensorDictBase, 選用) – 一個目標 tensordict。它必須具有與索引相同的形狀。

範例

>>> td = TensorDict(
...     {"a": torch.randn(3, 4, 5),
...      "b": TensorDict({"c": torch.zeros(3, 4, 5)}, [3, 4, 5])},
...     [3, 4])
>>> index = torch.randint(4, (3, 2))
>>> td_gather = td.gather(dim=1, index=index)
>>> print(td_gather)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 2, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 2]),
    device=None,
    is_shared=False)

Gather 保留維度名稱。

範例

>>> td.names = ["a", "b"]
>>> td_gather = td.gather(dim=1, index=index)
>>> td_gather.names
["a", "b"]

gather_and_stack(dst: int, group: 'dist.ProcessGroup' | None = None) → T | None¶

從各種 worker 收集 tensordict，並將它們堆疊到目標 worker 上的 self 中。

參數:

dst (int) – 將呼叫 gather_and_stack() 的目標 worker 的 rank。
group (torch.distributed.ProcessGroup, 選用) – 如果設定，指定的進程群組將用於通訊。否則，將使用預設的進程群組。預設為 None。

範例

>>> from torch import multiprocessing as mp
>>> from tensordict import TensorDict
>>> import torch
>>>
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...     # Create a single tensordict to be sent to server
...     td = TensorDict(
...         {("a", "b"): torch.randn(2),
...          "c": torch.randn(2)}, [2]
...     )
...     td.gather_and_stack(0)
...
>>> def server():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=0,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...     # Creates the destination tensordict on server.
...     # The first dim must be equal to world_size-1
...     td = TensorDict(
...         {("a", "b"): torch.zeros(2),
...          "c": torch.zeros(2)}, [2]
...     ).expand(1, 2).contiguous()
...     td.gather_and_stack(0)
...     assert td["a", "b"] != 0
...     print("yuppie")
...
>>> if __name__ == "__main__":
...     mp.set_start_method("spawn")
...
...     main_worker = mp.Process(target=server)
...     secondary_worker = mp.Process(target=client)
...
...     main_worker.start()
...     secondary_worker.start()
...
...     main_worker.join()
...     secondary_worker.join()

get(key: NestedKey, default: Any = None) → Tensor¶

取得使用輸入鍵儲存的值。

參數:

key (str, str 的元組) – 要查詢的鍵。如果是 str 的元組，則相當於鏈式呼叫 getattr。
default –
如果在 tensordict 中找不到該鍵，則為預設值。

警告

目前，如果 tensordict 中不存在鍵且未傳遞預設值，則會引發 KeyError。從 v0.7 開始，此行為將會改變，並且會改為傳回 None 值。若要採用新行為，請設定環境變數 export TD_GET_DEFAULTS_TO_NONE='1' 或呼叫 :func`~tensordict.set_get_defaults_to_none`。

範例

>>> td = TensorDict({"x": 1}, batch_size=[])
>>> td.get("x")
tensor(1)
>>> set_get_defaults_to_none(False) # Current default behaviour
>>> td.get("y") # Raises KeyError
>>> set_get_defaults_to_none(True)
>>> td.get("y")
None

get_at(key: NestedKey, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], default: Tensor = _NoDefault.ZERO) → Tensor¶

從 idx 索引處的鍵 key 取得 tensordict 的值。

參數:

key (str, str 的 tuple) – 要檢索的鍵。
index (int, slice, torch.Tensor, iterable) – tensor 的索引。
default (torch.Tensor) – 如果 tensordict 中不存在鍵，則傳回的預設值。

回傳值:

索引化的 tensor。

範例

>>> td = TensorDict({"x": torch.arange(3)}, batch_size=[])
>>> td.get_at("x", index=1)
tensor(1)

get_buffer(target: str) → Tensor¶

傳回 target 指定的緩衝區，如果不存在，則拋出錯誤。

有關此方法的功能以及如何正確指定 target 的更詳細說明，請參閱 get_submodule 的說明文件字串。

參數:: target – 要尋找的緩衝區的完整字串名稱。（有關如何指定完整字串，請參閱 get_submodule。）
回傳值:: 由 target 參考的緩衝區
回傳類型:: torch.Tensor
引發:: AttributeError – 如果目標字串參考了無效路徑或解析為非緩衝區的內容

get_extra_state() → Any¶

傳回任何要包含在模組的 state_dict 中的額外狀態。

如果您需要儲存額外狀態，請為您的模組實作此方法以及對應的 set_extra_state()。呼叫此函式會建置模組的 state_dict()。

請注意，額外狀態應可 pickle，以確保 state_dict 的序列化能正常運作。我們僅提供序列化 Tensor 的回溯相容性保證；如果其他物件的序列化 pickle 形式發生變更，則可能會破壞回溯相容性。

回傳值:: 要儲存在模組的 state_dict 中的任何額外狀態
回傳類型:: 物件

get_item_shape(key: NestedKey)¶: 傳回條目的形狀，可能會避免重新使用 get()。

get_non_tensor(key: NestedKey, default=_NoDefault.ZERO)¶

取得非 tensor 值（如果存在），如果找不到非 tensor 值，則取得 default。

此方法對於 tensor/TensorDict 值具備強大的適應性，這表示如果收集到的值是常規 tensor，也會傳回該值（雖然此方法會產生一些額外負擔，不應超出其自然範圍使用）。

有關如何在 tensordict 中設定非 tensor 值的更多資訊，請參閱 set_non_tensor()。

參數:

key (NestedKey) – NonTensorData 物件的位置。
default (Any, optional) – 如果找不到鍵，則傳回的值。

傳回值：tensordict.tensorclass.NonTensorData 的內容，: 或者如果 key 對應的條目不是 tensordict.tensorclass.NonTensorData，則返回該 key 對應的條目（如果找不到條目，則返回 default）。

範例

>>> data = TensorDict({}, batch_size=[])
>>> data.set_non_tensor(("nested", "the string"), "a string!")
>>> assert data.get_non_tensor(("nested", "the string")) == "a string!"
>>> # regular `get` works but returns a NonTensorData object
>>> data.get(("nested", "the string"))
NonTensorData(
    data='a string!',
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

get_parameter(target: str) → Parameter¶

如果存在 target 指定的參數，則返回該參數，否則拋出錯誤。

有關此方法的功能以及如何正確指定 target 的更詳細說明，請參閱 get_submodule 的說明文件字串。

參數:: target – 要查找的參數的完整字串名稱。（請參閱 get_submodule 以了解如何指定完整字串。）
回傳值:: 由 target 引用的參數
回傳類型:: torch.nn.Parameter
引發:: AttributeError – 如果目標字串引用無效的路徑，或解析為不是 nn.Parameter 的東西

get_submodule(target: str) → Module¶

如果存在 target 指定的子模組，則返回該子模組，否則拋出錯誤。

例如，假設您有一個 nn.Module A，如下所示

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

（該圖顯示了一個 nn.Module A。A 有一個巢狀的子模組 net_b，它本身有兩個子模組 net_c 和 linear。net_c 然後有一個子模組 conv。）

要檢查我們是否有 linear 子模組，我們會呼叫 get_submodule("net_b.linear")。要檢查我們是否有 conv 子模組，我們會呼叫 get_submodule("net_b.net_c.conv")。

get_submodule 的運行時間受 target 中模組巢狀結構的程度限制。針對 named_modules 的查詢可以達到相同的結果，但它是 O(N)，其中 N 是傳遞模組的數量。因此，對於簡單的檢查以查看某些子模組是否存在，應始終使用 get_submodule。

參數:: target – 要查找的子模組的完整字串名稱。（有關如何指定完整字串，請參閱上面的範例。）
回傳值:: 由 target 引用的子模組
回傳類型:: torch.nn.Module
引發:: AttributeError – 如果目標字串引用無效的路徑，或解析為不是 nn.Module 的東西

property grad¶: 返回一個包含葉張量的 .grad 屬性的 tensordict。

half()¶: 將所有張量轉換為 torch.half。

int()¶: 將所有張量轉換為 torch.int。

int16()¶: 將所有張量轉換為 torch.int16。

int32()¶: 將所有張量轉換為 torch.int32。

int64()¶: 將所有張量轉換為 torch.int64。

int8()¶: 將所有張量轉換為 torch.int8。

ipu(device: Optional[Union[int, device]] = None) → T¶

將所有模型參數和緩衝區移動到 IPU。

這也會使相關的參數和緩衝區成為不同的物件。因此，如果模組將在最佳化時駐留在 IPU 上，則應在建構最佳化器之前呼叫此方法。

注意

此方法會就地修改模組。

參數:: device (int, optional) – 如果指定，所有參數將被複製到該裝置
回傳值:: self
回傳類型:: 模組

irecv(src: int, *, group: 'dist.ProcessGroup' | None = None, return_premature: bool = False, init_tag: int = 0, pseudo_rand: bool = False) → tuple[int, list[torch.Future]] | list[torch.Future]] | None¶

非同步接收 tensordict 的內容，並使用其更新內容。

請查看 isend() 方法中的範例以了解上下文。

參數:

src (int) – 來源工作者的 rank。

關鍵字參數:

group (torch.distributed.ProcessGroup, 選用) – 如果設定，指定的進程群組將用於通訊。否則，將使用預設的進程群組。預設為 None。
return_premature (bool) – 如果為 True，則返回一個 future 列表，用於等待直到 tensordict 更新。預設為 False，即等待直到呼叫中完成更新。
init_tag (int) – 來源工作者使用的 init_tag。
pseudo_rand (bool) – 如果為 True，則標籤的序列將是偽隨機的，允許從不同節點發送多個數據而不會重疊。請注意，這些偽隨機數的生成成本很高（1e-5 秒/數字），這意味著它可能會減慢演算法的運行時間。此值必須與傳遞給 isend() 的值相符。預設為 False。

回傳值:

如果 return_premature=True，則為一個 futures 列表，用於等待: 直到更新 tensordict。

is_consolidated()¶: 檢查 TensorDict 是否具有 consolidated storage。

is_contiguous(*args, **kwargs)¶: 返回一個布林值，指示所有張量是否連續。

is_empty() → bool¶: 檢查 tensordict 是否包含任何 leaf。

property is_memmap: bool¶

檢查 tensordict 是否為 memory-mapped。

如果 TensorDict 實例是 memory-mapped，則它會被鎖定 (條目無法重新命名、刪除或新增)。如果使用都是 memory-mapped 的張量建立 TensorDict，這__並不__表示 is_memmap 將返回 True（因為新的張量可能是或可能不是 memory-mapped）。只有在呼叫 tensordict.memmap_() 時，tensordict 才會被視為 memory-mapped。

對於 CUDA 裝置上的 tensordict，這始終為 True。

property is_shared: bool¶

檢查 tensordict 是否位於共享記憶體中。

如果 TensorDict 實例位於共享記憶體中，它將被鎖定（條目無法重新命名、移除或新增）。如果使用位於共享記憶體中的張量建立 TensorDict，這__不__代表 is_shared 將返回 True（因為新的張量可能位於或不位於共享記憶體中）。只有當呼叫 tensordict.share_memory_() 或將 tensordict 放置在預設內容為共享的裝置上（例如，"cuda"），tensordict 才會被視為位於共享記憶體中。

對於 CUDA 裝置上的 tensordict，這始終為 True。

isend(dst: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → int¶

非同步發送 tensordict 的內容。

參數:

dst (int) – 內容應發送到的目標 worker 的排名。

關鍵字參數:

group (torch.distributed.ProcessGroup, 選用) – 如果設定，指定的進程群組將用於通訊。否則，將使用預設的進程群組。預設為 None。
init_tag (int) – 用於標記張量的初始標籤。請注意，這將根據 TensorDict 中包含的張量數量進行遞增。
pseudo_rand (bool) – 如果為 True，標籤的序列將是偽隨機的，允許從不同的節點發送多個資料而不會重疊。請注意，這些偽隨機數的產生成本很高（1e-5 秒/數字），這意味著它可能會降低演算法的運行時間。預設值為 False。

範例

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import multiprocessing as mp
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...
...     td = TensorDict(
...         {
...             ("a", "b"): torch.randn(2),
...             "c": torch.randn(2, 3),
...             "_": torch.ones(2, 1, 5),
...         },
...         [2],
...     )
...     td.isend(0)
...
>>>
>>> def server(queue, return_premature=True):
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=0,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...     td = TensorDict(
...         {
...             ("a", "b"): torch.zeros(2),
...             "c": torch.zeros(2, 3),
...             "_": torch.zeros(2, 1, 5),
...         },
...         [2],
...     )
...     out = td.irecv(1, return_premature=return_premature)
...     if return_premature:
...         for fut in out:
...             fut.wait()
...     assert (td != 0).all()
...     queue.put("yuppie")
...
>>>
>>> if __name__ == "__main__":
...     queue = mp.Queue(1)
...     main_worker = mp.Process(
...         target=server,
...         args=(queue, )
...         )
...     secondary_worker = mp.Process(target=client)
...
...     main_worker.start()
...     secondary_worker.start()
...     out = queue.get(timeout=10)
...     assert out == "yuppie"
...     main_worker.join()
...     secondary_worker.join()

isfinite() → T¶

返回一個新的 tensordict，其中包含布林元素，表示每個元素是否為有限值。

當實數值不是 NaN、負無限大或無限大時，它們是有限的。當複數的實部和虛部都是有限的時，它們是有限的。

isnan() → T¶

返回一個新的 tensordict，其中包含布林元素，表示輸入的每個元素是否為 NaN。

當複數的實部和/或虛部為 NaN 時，該複數被認為是 NaN。

isneginf() → T¶: 測試輸入的每個元素是否為負無限大。

isposinf() → T¶: 測試輸入的每個元素是否為負無限大。

isreal() → T¶: 返回一個新的 tensordict，其中包含布林元素，表示輸入的每個元素是否為實數值。

items(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) → Iterator[Tensor]¶

回傳 tensordict 的鍵值對 (key-value pair) 的產生器 (generator)。

參數:

include_nested (bool, optional) – 若為 True，則會回傳巢狀 (nested) 的值。預設為 False。
leaves_only (bool, optional) – 若為 False，則只會回傳葉節點 (leaves)。預設為 False。
is_leaf – 一個可選的可呼叫 (callable) 物件，用於指示一個類別是否應被視為葉節點。

關鍵字參數:

sort (bool, optional) – 是否應該排序鍵。對於巢狀鍵，鍵會根據它們的連接名稱進行排序 (例如，("a", "key") 將被計為 "a.key" 以進行排序)。請注意，在處理大型 tensordict 時，排序可能會產生顯著的開銷。預設為 False。

keys(*args, **kwargs)¶

回傳 tensordict 鍵的產生器。

警告

TensorDict 的 keys() 方法回傳鍵的惰性視圖 (lazy view)。如果查詢了 keys 但未迭代它們，然後修改了 tensordict，稍後迭代這些鍵將會回傳鍵的新配置。

參數:

include_nested (bool, optional) – 若為 True，則會回傳巢狀 (nested) 的值。預設為 False。
leaves_only (bool, optional) – 若為 False，則只會回傳葉節點 (leaves)。預設為 False。
is_leaf – 一個可選的可呼叫 (callable) 物件，用於指示一個類別是否應被視為葉節點。

關鍵字參數:

sort (bool, optional) – 鍵是否應該排序。對於巢狀鍵，鍵會根據它們的連接名稱進行排序 (例如，("a", "key") 將被計為 "a.key" 以進行排序)。請注意，在處理大型 tensordict 時，排序可能會產生顯著的開銷。預設為 False。

範例

>>> from tensordict import TensorDict
>>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[])
>>> data.keys()
['0', '1']
>>> list(data.keys(leaves_only=True))
['0']
>>> list(data.keys(include_nested=True, leaves_only=True))
['0', '1', ('1', '2')]

classmethod lazy_stack(input, dim=0, *, out=None, **kwargs)¶

建立 tensordict 的惰性堆疊 (lazy stack)。

詳情請參閱 lazy_stack()。

lerp(end: tensordict.base.TensorDictBase | torch.Tensor, weight: tensordict.base.TensorDictBase | torch.Tensor | float)¶

基於純量或張量 weight，對兩個張量 start (由 self 給定) 和 end 進行線性內插 (linear interpolation)。

\[\text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i)\]

start 和 end 的形狀必須是可廣播的 (broadcastable)。如果 weight 是一個張量，那麼 weight、start 和 end 的形狀必須是可廣播的。

參數:

end (TensorDict) – 具有結束點 (ending points) 的 tensordict。
weight (TensorDict, tensor 或 float) – 用於內插公式的權重。

lerp_(end: tensordict.base.TensorDictBase | float, weight: tensordict.base.TensorDictBase | float)¶: lerp() 的原地 (in-place) 版本。

lgamma() → T¶: 計算 TensorDict 中每個元素的 lgamma() 值。

lgamma_() → T¶: 計算 TensorDict 中每個元素的 lgamma() 值，並以原地 (in-place) 方式進行。

classmethod load(prefix: str | pathlib.Path, *args, **kwargs) → T¶

從磁碟載入 tensordict。

這個類別方法是 load_memmap() 的代理。

load_(prefix: str | pathlib.Path, *args, **kwargs)¶

在目前的 tensordict 中，從磁碟載入 tensordict。

這個類別方法是 load_memmap_() 的代理。

classmethod load_memmap(prefix: str | pathlib.Path, device: Optional[device] = None, non_blocking: bool = False, *, out: Optional[TensorDictBase] = None) → T¶

從磁碟載入記憶體映射 (memory-mapped) 的 tensordict。

參數:

prefix (str 或 資料夾路徑) – 儲存的 tensordict 應從該路徑的資料夾中提取。
device (torch.device 或 等效裝置, optional) – 如果提供，資料將會非同步地轉換到該裝置。支援 “meta” 裝置，在這種情況下，資料不會被載入，而是創建一組空的 “meta” tensors。這對於了解總體模型大小和結構，而無需實際打開任何檔案非常有用。
non_blocking (bool, optional) – 如果 True，在將 tensors 載入裝置後，將不會調用同步 (synchronize)。預設值為 False。
out (TensorDictBase, optional) – 可選的 tensordict，資料應寫入其中。

範例

>>> from tensordict import TensorDict
>>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0)
>>> td.memmap("./saved_td")
>>> td_load = TensorDict.load_memmap("./saved_td")
>>> assert (td == td_load).all()

此方法還允許載入巢狀的 tensordicts。

範例

>>> nested = TensorDict.load_memmap("./saved_td/nested")
>>> assert nested["e"] == 0

tensordict 也可以載入到 "meta" 裝置上，或者作為一個假的 tensor。

範例

>>> import tempfile
>>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> with tempfile.TemporaryDirectory() as path:
...     td.save(path)
...     td_load = TensorDict.load_memmap(path, device="meta")
...     print("meta:", td_load)
...     from torch._subclasses import FakeTensorMode
...     with FakeTensorMode():
...         td_load = TensorDict.load_memmap(path)
...         print("fake:", td_load)
meta: TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=meta,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=meta,
    is_shared=False)
fake: TensorDict(
    fields={
        a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

load_memmap_(prefix: str | pathlib.Path)¶

將記憶體映射 tensordict 的內容載入到呼叫 load_memmap_ 的 tensordict 中。

請參閱 load_memmap() 以取得更多資訊。

load_state_dict(state_dict: OrderedDict[str, Any], strict=True, assign=False)¶

載入一個 state_dict，其格式如同 state_dict()，到 tensordict 中。

參數:

state_dict (OrderedDict) – 要複製的 state_dict。
strict (bool, optional) – 是否嚴格強制 state_dict 中的鍵與此 tensordict 的 torch.nn.Module.state_dict() 函數傳回的鍵匹配。預設值：True
assign (bool, optional) – 是否將 state dictionary 中的項目賦值到 tensordict 中對應的鍵，而不是將它們就地複製到 tensordict 目前的 tensors 中。當 False 時，目前模組中的 tensors 的屬性會被保留，而當 True 時，state dict 中的 Tensors 的屬性會被保留。預設值：False
from_flatten (bool, optional) – 如果 True，則假設輸入的 state_dict 已被扁平化。預設為 False。

範例

>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, [])
>>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, [])
>>> sd = data.state_dict()
>>> data_zeroed.load_state_dict(sd)
>>> print(data_zeroed["3", "3"])
tensor(3)
>>> # with flattening
>>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, [])
>>> data_zeroed.load_state_dict(data.state_dict(flatten=True), from_flatten=True)
>>> print(data_zeroed["3", "3"])
tensor(3)

lock_() → T¶

鎖定一個 tensordict，使其無法進行非就地操作。

諸如 set()、__setitem__()、update()、rename_key_() 或其他新增或移除條目的操作將會被阻止。

此方法可以用作裝飾器。

範例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 1, "b": 2, "c": 3}, batch_size=[])
>>> with td.lock_():
...     assert td.is_locked
...     try:
...         td.set("d", 0) # error!
...     except RuntimeError:
...         print("td is locked!")
...     try:
...         del td["d"]
...     except RuntimeError:
...         print("td is locked!")
...     try:
...         td.rename_key_("a", "d")
...     except RuntimeError:
...         print("td is locked!")
...     td.set("a", 0, inplace=True)  # No storage is added, moved or removed
...     td.set_("a", 0) # No storage is added, moved or removed
...     td.update({"a": 0}, inplace=True)  # No storage is added, moved or removed
...     td.update_({"a": 0})  # No storage is added, moved or removed
>>> assert not td.is_locked

log() → T¶: 計算 TensorDict 中每個元素的 log() 值。

log10() → T¶: 計算 TensorDict 中每個元素的 log10() 值。

log10_() → T¶: 就地計算 TensorDict 中每個元素的 log10() 值。

log1p() → T¶: 計算 TensorDict 中每個元素的 log1p() 值。

log1p_() → T¶: 就地計算 TensorDict 中每個元素的 log1p() 值。

log2() → T¶: 計算 TensorDict 中每個元素的 log2() 值。

log2_() → T¶: 就地計算 TensorDict 中每個元素的 log2() 值。

log_() → T¶: 就地計算 TensorDict 中每個元素的 log() 值。

make_memmap(key: NestedKey, shape: torch.Size | torch.Tensor, *, dtype: Optional[dtype] = None) → MemoryMappedTensor¶

根據形狀（shape）以及可能的資料類型（dtype）建立一個空的記憶體映射張量（memory-mapped tensor）。

警告

此方法在設計上並非鎖定安全（lock-safe）的。在多個節點上存在的記憶體映射 TensorDict 實例將需要使用 memmap_refresh_() 方法來更新。

寫入現有的條目將會導致錯誤。

參數:

key (NestedKey) – 要寫入的新條目的鍵（key）。如果該鍵已經存在於 tensordict 中，則會引發例外。
shape (torch.Size 或等效值, torch.Tensor 用於巢狀張量) – 要寫入的張量的形狀。

關鍵字參數:

dtype (torch.dtype, optional) – 新張量的資料類型。

回傳值:

一個新的記憶體映射張量。

make_memmap_from_storage(key: NestedKey, storage: UntypedStorage, shape: torch.Size | torch.Tensor, *, dtype: Optional[dtype] = None) → MemoryMappedTensor¶

根據儲存體（storage）、形狀以及可能的資料類型建立一個空的記憶體映射張量。

警告

此方法在設計上並非鎖定安全（lock-safe）的。在多個節點上存在的記憶體映射 TensorDict 實例將需要使用 memmap_refresh_() 方法來更新。

注意

如果儲存體有關聯的檔名，它必須與檔案的新檔名相符。如果它沒有關聯的檔名，但 tensordict 有關聯的路徑，這將導致例外。

參數:

key (NestedKey) – 要寫入的新條目的鍵（key）。如果該鍵已經存在於 tensordict 中，則會引發例外。
storage (torch.UntypedStorage) – 用於新 MemoryMappedTensor 的儲存體。必須是物理記憶體儲存體。
shape (torch.Size 或等效值, torch.Tensor 用於巢狀張量) – 要寫入的張量的形狀。

關鍵字參數:

dtype (torch.dtype, optional) – 新張量的資料類型。

回傳值:

一個具有給定儲存體的新記憶體映射張量。

make_memmap_from_tensor(key: NestedKey, tensor: Tensor, *, copy_data: bool = True) → MemoryMappedTensor¶

根據一個張量（tensor）建立一個空的記憶體映射張量。

警告

此方法在設計上並非鎖定安全（lock-safe）的。在多個節點上存在的記憶體映射 TensorDict 實例將需要使用 memmap_refresh_() 方法來更新。

如果 copy_data 為 True，此方法總是複製儲存體內容（即，儲存體不共享）。

參數:

key (NestedKey) – 要寫入的新條目的鍵（key）。如果該鍵已經存在於 tensordict 中，則會引發例外。
tensor (torch.Tensor) – 要在物理記憶體上複製的張量。

關鍵字參數:

copy_data (bool, optional) – 如果 False，則新張量將共享輸入的元數據（例如形狀和資料類型），但內容將為空。預設值為 True。

回傳值:

一個具有給定儲存體的新記憶體映射張量。

map(fn: Callable, dim: int = 0, num_workers: Optional[int] = None, chunksize: Optional[int] = None, num_chunks: Optional[int] = None, pool: Optional[Pool] = None, generator: Optional[Generator] = None, max_tasks_per_child: Optional[int] = None, worker_threads: int = 1, mp_start_method: Optional[str] = None)¶

將函式映射到 tensordict 在一個維度上的分割。

此方法會將函式應用於 tensordict 實例，方法是將其分塊為大小相等的 tensordict，並將操作分派到所需數量的 worker 上。

函式簽章應為 Callabe[[TensorDict], Union[TensorDict, Tensor]]。輸出必須支援 torch.cat() 操作。該函式必須可序列化。

注意

此方法在處理儲存在磁碟上的大型資料集 (例如，記憶體對應的 tensordict) 時特別有用，在這些資料集中，chunks 將是原始資料的零複製切片，可以以幾乎零成本的方式傳遞到各個程序。這允許以很小的成本處理非常大的資料集 (例如，超過 Tb 大小)。

參數:

fn (callable) – 應用於 tensordict 的函式。支援類似於 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的簽章。
dim (int, optional) – tensordict 將被分塊的維度。
num_workers (int, optional) – worker 的數量。與 pool 互斥。如果未提供，worker 的數量將設定為可用的 cpu 數量。

關鍵字參數:

out (TensorDictBase, optional) – 輸出的可選容器。其沿著提供的 dim 維度的批次大小必須與 self.ndim 匹配。如果它是共享的或記憶體映射的 ( is_shared() 或 is_memmap() 返回 True) ，它將在遠端程序中填充，避免資料傳輸。否則，來自 self 切片的資料將被發送到該程序，在當前程序中收集，並原地寫入 out。
chunksize (int, optional) – 每個資料塊的大小。chunksize 為 0 將沿著所需的維度解綁 tensordict，並在應用該函數後重新堆疊它，而 chunksize>0 將分割 tensordict 並在產生的 tensordicts 列表中調用 torch.cat()。如果未提供，則區塊的數量將等於 worker 的數量。對於非常大的 tensordicts，這麼大的區塊可能不適合記憶體來完成操作，並且可能需要更多區塊才能使操作實際可行。此參數與 num_chunks 互斥。
num_chunks (int, optional) – 將 tensordict 分割成的區塊數。如果未提供，則區塊的數量將等於 worker 的數量。對於非常大的 tensordicts，這麼大的區塊可能不適合記憶體來完成操作，並且可能需要更多區塊才能使操作實際可行。此參數與 chunksize 互斥。
pool (mp.Pool, optional) – 用於執行任務的多處理 Pool 實例。如果未提供，則將在 map 方法中建立一個 pool。
generator (torch.Generator, optional) –
用於設定種子的產生器。將從中產生一個基本種子，並且 pool 的每個 worker 將以提供的種子遞增一個從 0 到 num_workers 的唯一整數來設定種子。如果未提供產生器，則將使用隨機整數作為種子。要使用未設定種子的 worker，應單獨建立 pool 並直接傳遞到 map()。

注意

提供低值種子時應謹慎，因為這可能會導致實驗之間的自相關，例如：如果要求 8 個 worker 並且種子為 4，則 worker 的種子範圍將從 4 到 11。如果種子為 5，則 worker 的種子範圍將從 5 到 12。這兩個實驗將有 7 個種子的重疊，這可能會對結果產生意想不到的影響。

注意

worker 設定種子的目的是在每個 worker 上擁有獨立的種子，而不是在 map 方法的調用之間產生可重複的結果。換句話說，兩個實驗可能會並且可能會返回不同的結果，因為無法知道哪個 worker 將選擇哪個任務。但是，我們可以確保每個 worker 都有不同的種子，並且每個 worker 上的偽隨機操作將是不相關的。
max_tasks_per_child (int, optional) – 每個子程序選取的最大任務數。預設為 None，即對任務數沒有限制。
worker_threads (int, optional) – worker 的線程數。預設為 1。
index_with_generator (bool, optional) – 如果 True，則在查詢期間完成 tensordict 的分割/分塊，從而節省初始化時間。請注意，chunk() 和 split() 比索引（在產生器中使用）效率更高，因此在初始化時節省處理時間可能會對總運行時間產生負面影響。預設為 False。
pbar (bool, optional) – 如果 True，將顯示進度條。需要提供 tqdm。預設為 False。
mp_start_method (str, optional) – 多處理的啟動方法。如果未提供，將使用預設啟動方法。接受的字串為 "fork" 和 "spawn"。請記住，"cuda" 張量無法在使用 "fork" 啟動方法的程序之間共享。如果將 pool 傳遞給 map 方法，則此操作無效。

範例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> def process_data(data):
...     data.set("y", data.get("x") + 1)
...     return data
>>> if __name__ == "__main__":
...     data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_()
...     data = data.map(process_data, dim=1)
...     print(data["y"][:, :10])
...
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

map_iter(fn: Callable[[TensorDictBase], TensorDictBase | None], dim: int = 0, num_workers: int | None = None, *, shuffle: bool = False, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: torch.Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = True, pbar: bool = False, mp_start_method: str | None = None)¶

迭代地將函式映射到 tensordict 沿一個維度的分割。

這是 map() 的可迭代版本。

此方法將通過將 tensordict 實例分成大小相等的 tensordict，並在所需數量的 worker 上分派操作，從而將函式應用於 tensordict 實例。它將一次產生一個結果。

函式簽名應為 Callabe[[TensorDict], Union[TensorDict, Tensor]]。該函式必須是可序列化的。

注意

此方法在處理儲存在磁碟上的大型資料集 (例如，記憶體對應的 tensordict) 時特別有用，在這些資料集中，chunks 將是原始資料的零複製切片，可以以幾乎零成本的方式傳遞到各個程序。這允許以很小的成本處理非常大的資料集 (例如，超過 Tb 大小)。

注意

此函式可用於表示資料集並從中載入，以類似 dataloader 的方式進行。

參數:

fn (callable) – 應用於 tensordict 的函式。支援類似於 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的簽章。
dim (int, optional) – tensordict 將被分塊的維度。
num_workers (int, optional) – worker 的數量。與 pool 互斥。如果未提供，worker 的數量將設定為可用的 cpu 數量。

關鍵字參數:

shuffle (bool, optional) – 是否應對索引進行全域洗牌。如果 True，則每個批次將包含非連續的樣本。如果 index_with_generator=False 且 shuffle=True`，則會引發錯誤。預設值為 False。
chunksize (int, optional) – 每個資料塊的大小。chunksize 為 0 將沿著所需的維度解綁 tensordict，並在應用該函數後重新堆疊它，而 chunksize>0 將分割 tensordict 並在產生的 tensordicts 列表中調用 torch.cat()。如果未提供，則區塊的數量將等於 worker 的數量。對於非常大的 tensordicts，這麼大的區塊可能不適合記憶體來完成操作，並且可能需要更多區塊才能使操作實際可行。此參數與 num_chunks 互斥。
num_chunks (int, optional) – 將 tensordict 分割成的區塊數。如果未提供，則區塊的數量將等於 worker 的數量。對於非常大的 tensordicts，這麼大的區塊可能不適合記憶體來完成操作，並且可能需要更多區塊才能使操作實際可行。此參數與 chunksize 互斥。
pool (mp.Pool, optional) – 用於執行任務的多處理 Pool 實例。如果未提供，則將在 map 方法中建立一個 pool。
generator (torch.Generator, optional) –
用於設定種子的產生器。將從中產生一個基本種子，並且 pool 的每個 worker 將以提供的種子遞增一個從 0 到 num_workers 的唯一整數來設定種子。如果未提供產生器，則將使用隨機整數作為種子。要使用未設定種子的 worker，應單獨建立 pool 並直接傳遞到 map()。

注意

提供低值種子時應謹慎，因為這可能會導致實驗之間的自相關，例如：如果要求 8 個 worker 並且種子為 4，則 worker 的種子範圍將從 4 到 11。如果種子為 5，則 worker 的種子範圍將從 5 到 12。這兩個實驗將有 7 個種子的重疊，這可能會對結果產生意想不到的影響。

注意

worker 設定種子的目的是在每個 worker 上擁有獨立的種子，而不是在 map 方法的調用之間產生可重複的結果。換句話說，兩個實驗可能會並且可能會返回不同的結果，因為無法知道哪個 worker 將選擇哪個任務。但是，我們可以確保每個 worker 都有不同的種子，並且每個 worker 上的偽隨機操作將是不相關的。
max_tasks_per_child (int, optional) – 每個子程序選取的最大任務數。預設為 None，即對任務數沒有限制。
worker_threads (int, optional) – worker 的線程數。預設為 1。
index_with_generator (bool, optional) –
如果 True，則 tensordict 的分割/分塊將在查詢期間完成，從而節省初始化時間。請注意，chunk() 和 split() 比索引（在生成器中使用）有效率得多，因此在初始化時間節省處理時間可能會對總執行時間產生負面影響。預設值為 True。

注意

index_with_generator 的預設值對於 map_iter 和 map 而言有所不同，前者假設在記憶體中儲存分割版本的 TensorDict 成本非常高。
pbar (bool, optional) – 如果 True，將顯示進度條。需要提供 tqdm。預設為 False。
mp_start_method (str, optional) – 多處理的啟動方法。如果未提供，將使用預設啟動方法。接受的字串為 "fork" 和 "spawn"。請記住，"cuda" 張量無法在使用 "fork" 啟動方法的程序之間共享。如果將 pool 傳遞給 map 方法，則此操作無效。

範例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> def process_data(data):
...     data.unlock_()
...     data.set("y", data.get("x") + 1)
...     return data
>>> if __name__ == "__main__":
...     data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_()
...     for sample in data.map_iter(process_data, dim=1, chunksize=5):
...         print(sample["y"])
...         break
...
tensor([[1., 1., 1., 1., 1.]])

masked_fill(*args, **kwargs)¶

masked_fill 的異地 (out-of-place) 版本。

參數:

mask (boolean torch.Tensor) – 要填充的值的遮罩。形狀必須與 tensordict 的批次大小 (batch-size) 相符。
value – 用於填充張量的值。

回傳值:

self

範例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)},
...     batch_size=[3])
>>> mask = torch.tensor([True, False, False])
>>> td1 = td.masked_fill(mask, 1.0)
>>> td1.get("a")
tensor([[1., 1., 1., 1.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])

masked_fill_(*args, **kwargs)¶

使用所需的值填充對應於遮罩的值。

參數:

mask (boolean torch.Tensor) – 要填充的值的遮罩。形狀必須與 tensordict 的批次大小 (batch-size) 相符。
value – 用於填充張量的值。

回傳值:

self

範例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)},
...     batch_size=[3])
>>> mask = torch.tensor([True, False, False])
>>> td.masked_fill_(mask, 1.0)
>>> td.get("a")
tensor([[1., 1., 1., 1.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])

masked_select(mask: Tensor) → T¶

遮罩 TensorDict 的所有張量，並傳回一個新的 TensorDict 實例，該實例具有指向遮罩值的相似鍵。

參數:: mask (torch.Tensor) – 用於張量的布林遮罩。形狀必須與 TensorDict 的 batch_size 相符。

範例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)},
...    batch_size=[3])
>>> mask = torch.tensor([True, False, False])
>>> td_mask = td.masked_select(mask)
>>> td_mask.get("a")
tensor([[0., 0., 0., 0.]])

maximum(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T¶

計算 self 和 other 的元素級別的最大值。

參數:: other (TensorDict or Tensor) – 另一個輸入 tensordict 或 tensor。
關鍵字參數:: default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

maximum_(other: tensordict.base.TensorDictBase | torch.Tensor) → T¶: maximum() 的原地 (in-place) 版本。

注意

原地 maximum 不支援 default 關鍵字引數。

classmethod maybe_dense_stack(input, dim=0, *, out=None, **kwargs)¶

嘗試建立 tensordict 的密集堆疊，並在需要時回退到延遲堆疊。

有關詳細資訊，請參閱 maybe_dense_stack()。

mean(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor¶

傳回輸入 tensordict 中所有元素的平均值。

參數:

dim (int, int 的 tuple, optional) – 如果 None，則傳回包含所有 leaves 平均值 (如果可以計算) 的無維度 tensordict。如果是整數或整數的 tuple，則只有當此維度與 tensordict 形狀相容時，才會在指定的維度上呼叫 mean。
keepdim (bool) – 輸出 tensor 是否保留 dim。

關鍵字參數:

dtype (torch.dtype, optional) – 傳回的 tensor 所需的資料類型。如果指定，則在執行操作之前，輸入 tensor 會被轉換為 dtype。這對於防止資料類型溢位很有用。預設值： None。
reduce (bool, optional) – 如果 True，則 reduciton 將發生在所有 TensorDict 值上，並傳回單個簡化的 tensor。預設為 False。

memmap(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) → T¶

將所有 tensors 寫入新的 tensordict 中對應的記憶體對應 Tensor。

參數:

prefix (str) – 記憶體對應 tensors 將儲存的目錄前綴。目錄樹狀結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False（預設值），如果 tensordict 中的條目已經是一個儲存在磁碟上且具有關聯檔案的 tensor，但未根據 prefix 儲存在正確的位置，則會引發例外。如果為 True，任何現有的 Tensor 都將被複製到新的位置。

關鍵字參數:

num_threads (int, optional) – 用於寫入 memmap tensors 的執行緒數量。預設為 0。
return_early (bool, optional) – 若 True 且 num_threads>0，此方法將會回傳 tensordict 的 future。
share_non_tensor (bool, optional) – 若 True，非 tensor 的資料將會在各個 process 之間共享，而且在單一節點內任何 worker 上進行的寫入操作（例如 inplace 更新或設定）都會更新所有其他 worker 上的值。如果非 tensor 的 leaves 數量很高（例如，共享大量的非 tensor 資料），這可能會導致記憶體不足（OOM）或類似的錯誤。預設值為 False。
existsok (bool, optional) – 若 False，如果 tensor 已經存在於相同的路徑中，將會引發例外。預設值為 True。

然後 TensorDict 會被鎖定，代表任何非 in-place 的寫入操作都會拋出例外（例如，重新命名、設定或移除一個 entry）。一旦 tensordict 被解鎖，memory-mapped 屬性就會變成 False，因為跨 process 的 identity 不再保證。

回傳值:: 如果 return_early=False，則會產生一個新的 tensordict，其中 tensor 儲存在磁碟上，否則會產生一個 TensorDictFuture 實例。

注意

以這種方式進行序列化對於深度巢狀的 tensordict 來說可能會很慢，因此不建議在訓練迴圈中呼叫此方法。

memmap_(prefix: Optional[str] = None, copy_existing: bool = False, num_threads: int = 0) → TensorDictBase¶

將所有 tensor in-place 寫入到對應的 memory-mapped Tensor 上。

參數:

prefix (str) – 記憶體對應 tensors 將儲存的目錄前綴。目錄樹狀結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False（預設值），如果 tensordict 中的條目已經是一個儲存在磁碟上且具有關聯檔案的 tensor，但未根據 prefix 儲存在正確的位置，則會引發例外。如果為 True，任何現有的 Tensor 都將被複製到新的位置。

關鍵字參數:

num_threads (int, optional) – 用於寫入 memmap tensors 的執行緒數量。預設為 0。
return_early (bool, optional) – 如果 True 且 num_threads>0，則該方法將傳回 tensordict 的 future。可以使用 future.result() 查詢產生的 tensordict。
share_non_tensor (bool, optional) – 若 True，非 tensor 的資料將會在各個 process 之間共享，而且在單一節點內任何 worker 上進行的寫入操作（例如 inplace 更新或設定）都會更新所有其他 worker 上的值。如果非 tensor 的 leaves 數量很高（例如，共享大量的非 tensor 資料），這可能會導致記憶體不足（OOM）或類似的錯誤。預設值為 False。
existsok (bool, optional) – 若 False，如果 tensor 已經存在於相同的路徑中，將會引發例外。預設值為 True。

然後 TensorDict 會被鎖定，代表任何非 in-place 的寫入操作都會拋出例外（例如，重新命名、設定或移除一個 entry）。一旦 tensordict 被解鎖，memory-mapped 屬性就會變成 False，因為跨 process 的 identity 不再保證。

回傳值:: 如果 return_early=False，則回傳 self，否則回傳 TensorDictFuture 實例。

注意

以這種方式進行序列化對於深度巢狀的 tensordict 來說可能會很慢，因此不建議在訓練迴圈中呼叫此方法。

memmap_like(prefix: str | None = None, copy_existing: bool = False, num_threads: int = 0) → T¶

建立一個無內容的 Memory-mapped tensordict，其形狀與原始 tensordict 相同。

參數:

prefix (str) – 記憶體對應 tensors 將儲存的目錄前綴。目錄樹狀結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False（預設值），如果 tensordict 中的條目已經是一個儲存在磁碟上且具有關聯檔案的 tensor，但未根據 prefix 儲存在正確的位置，則會引發例外。如果為 True，任何現有的 Tensor 都將被複製到新的位置。

關鍵字參數:

num_threads (int, optional) – 用於寫入 memmap tensors 的執行緒數量。預設為 0。
return_early (bool, optional) – 若 True 且 num_threads>0，此方法將會回傳 tensordict 的 future。
share_non_tensor (bool, optional) – 若 True，非 tensor 的資料將會在各個 process 之間共享，而且在單一節點內任何 worker 上進行的寫入操作（例如 inplace 更新或設定）都會更新所有其他 worker 上的值。如果非 tensor 的 leaves 數量很高（例如，共享大量的非 tensor 資料），這可能會導致記憶體不足（OOM）或類似的錯誤。預設值為 False。
existsok (bool, optional) – 若 False，如果 tensor 已經存在於相同的路徑中，將會引發例外。預設值為 True。

然後 TensorDict 會被鎖定，代表任何非 in-place 的寫入操作都會拋出例外（例如，重新命名、設定或移除一個 entry）。一旦 tensordict 被解鎖，memory-mapped 屬性就會變成 False，因為跨 process 的 identity 不再保證。

回傳值:: 如果 return_early=False，則會產生一個新的 TensorDict 實例，其中資料儲存為 memory-mapped tensor，否則會產生一個 TensorDictFuture 實例。

注意

這是將一組大型 buffer 寫入磁碟的建議方法，因為 memmap_() 將複製資訊，對於大型內容而言，這可能會很慢。

範例

>>> td = TensorDict({
...     "a": torch.zeros((3, 64, 64), dtype=torch.uint8),
...     "b": torch.zeros(1, dtype=torch.int64),
... }, batch_size=[]).expand(1_000_000)  # expand does not allocate new memory
>>> buffer = td.memmap_like("/path/to/dataset")

memmap_refresh_()¶

如果 memory-mapped tensordict 具有 saved_path，則會重新整理其內容。

如果沒有與其關聯的路徑，此方法將引發例外。

minimum(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T¶

計算 self 和 other 的 element-wise 最小值。

參數:: other (TensorDict or Tensor) – 另一個輸入 tensordict 或 tensor。
關鍵字參數:: default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

minimum_(other: tensordict.base.TensorDictBase | torch.Tensor) → T¶: minimum() 的原地 (in-place) 版本。

注意

原地 minimum 不支援 default 關鍵字參數。

modules() → Iterator[Module]¶

傳回網路中所有模組的迭代器。

產生：: Module – 網路中的一個模組

注意

重複的模組只會傳回一次。在以下範例中，l 只會傳回一次。

範例

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

mtia(device: Optional[Union[int, device]] = None) → T¶

將所有模型參數和緩衝區移動到 MTIA。

這也會使關聯的參數和緩衝區成為不同的物件。因此，如果模組將在 MTIA 上存在並進行最佳化，則應在建構最佳化器之前呼叫此方法。

注意

此方法會就地修改模組。

參數:: device (int, optional) – 如果指定，所有參數將被複製到該裝置
回傳值:: self
回傳類型:: 模組

mul(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T¶

將 other 乘以 self。

\[\text{{out}}_i = \text{{input}}_i \times \text{{other}}_i\]

支援廣播 (broadcasting)、類型提升 (type promotion)，以及整數、浮點數和複數輸入。

參數:: other (TensorDict, Tensor 或 Number) – 要從 self 中減去的張量或數字。
關鍵字參數:: default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

mul_(other: tensordict.base.TensorDictBase | torch.Tensor) → T¶: mul() 的原地 (in-place) 版本。

注意

原地 mul 不支援 default 關鍵字參數。

named_apply(fn: Callable, *others: TensorDictBase, batch_size: Optional[Sequence[int]] = None, device: torch.device | None = _NoDefault.ZERO, names: Optional[Sequence[str]] = _NoDefault.ZERO, inplace: bool = False, default: Any = _NoDefault.ZERO, filter_empty: bool | None = None, call_on_nested: bool = False, **constructor_kwargs) → tensordict.base.TensorDictBase | None¶

將鍵值條件可調用的函式應用於 tensordict 中儲存的所有值，並將它們設定在新的 atensordict 中。

可調用的函式簽章必須是 Callable[Tuple[str, Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。

參數:

fn (Callable) – 要應用於 tensordict 中 (name, tensor) 配對的函式。對於每個葉節點，只會使用它的葉節點名稱 (而不是完整的 NestedKey)。
*others (TensorDictBase 實例, optional) – 如果提供，這些 tensordict 實例應該具有與 self 相同的結構。fn 參數應該接收與 tensordict 數量一樣多的未命名輸入，包括 self。如果其他 tensordict 缺少條目，則可以透過 default 關鍵字參數傳遞預設值。
nested_keys (bool, optional) – 如果 True，則將使用葉節點的完整路徑。預設為 False，也就是說，只有最後一個字串會傳遞給函式。
batch_size (sequence of int, optional) – 如果提供此參數，產生的 TensorDict 將具有期望的 batch_size。 batch_size 參數應與轉換後的 batch_size 相符。這是一個僅限關鍵字的參數。
device (torch.device, optional) – 產生的 device，如果有的話。
names (list of str, optional) – 新的維度名稱，在 batch_size 被修改的情況下。
inplace (bool, optional) – 如果為 True，則會就地進行變更。預設值為 False。這是一個僅限關鍵字的引數。
default (Any, optional) – 其他 tensordict 中缺失條目的預設值。如果未提供，缺失的條目將引發 KeyError。
filter_empty (bool, optional) – 如果 True，則將過濾掉空的 tensordict。這也具有較低的計算成本，因為不會建立和銷毀空的資料結構。為了向後相容，預設值為 False。
propagate_lock (bool, optional) – 如果 True，則鎖定的 tensordict 將產生另一個鎖定的 tensordict。預設值為 False。

call_on_nested (bool, optional) –

如果 True，該函數將在第一層張量和容器（TensorDict 或 tensorclass）上呼叫。在這種情況下，func 負責將其呼叫傳播到巢狀層級。這允許在將呼叫傳播到巢狀 tensordict 時進行細粒度的行為控制。如果 False，該函數只會在葉節點上呼叫，並且 apply 會負責將該函數分派到所有葉節點。

>>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]})
>>> def mean_tensor_only(val):
...     if is_tensor_collection(val):
...         raise RuntimeError("Unexpected!")
...     return val.mean()
>>> td_mean = td.apply(mean_tensor_only)
>>> def mean_any(val):
...     if is_tensor_collection(val):
...         # Recurse
...         return val.apply(mean_any, call_on_nested=True)
...     return val.mean()
>>> td_mean = td.apply(mean_any, call_on_nested=True)

out (TensorDictBase, optional) –
一個用於寫入結果的 tensordict。這可以用來避免創建新的 tensordict。
```
>>> td = TensorDict({"a": 0})
>>> td.apply(lambda x: x+1, out=td)
>>> assert (td==1).all()
```
警告

如果在 tensordict 上執行的操作需要訪問多個鍵才能進行單一計算，則提供一個等於 self 的 out 參數可能會導致該操作默默地提供錯誤的結果。例如
```
>>> td = TensorDict({"a": 1, "b": 1})
>>> td.apply(lambda x: x+td["a"])["b"] # Right!
tensor(2)
>>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong!
tensor(3)
```
**constructor_kwargs – 要傳遞給 TensorDict 建構子的其他關鍵字參數。

回傳值:

一個帶有 transformed_in 張量的新 tensordict。

範例

>>> td = TensorDict({
...     "a": -torch.ones(3),
...     "nested": {"a": torch.ones(3), "b": torch.zeros(3)}},
...     batch_size=[3])
>>> def name_filter(name, tensor):
...     if name == "a":
...         return tensor
>>> td.named_apply(name_filter)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> def name_filter(name, *tensors):
...     if name == "a":
...         r = 0
...         for tensor in tensors:
...             r = r + tensor
...         return tensor
>>> out = td.named_apply(name_filter, td)
>>> print(out)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> print(out["a"])
tensor([-1., -1., -1.])

注意

如果函數回傳 None，則會忽略該條目。這可以用來過濾 tensordict 中的資料。

>>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, [])
>>> def name_filter(name, tensor):
...     if name == "1":
...         return tensor
>>> td.named_apply(name_filter)
TensorDict(
    fields={
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

named_buffers(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) → Iterator[Tuple[str, Tensor]]¶

回傳一個疊代器，遍歷模組緩衝區 (buffer)，產生緩衝區的名稱和緩衝區本身。

參數:

prefix (str) – 要加在所有緩衝區名稱前面的前綴詞。
recurse (bool, optional) – 如果為 True，則產生此模組和所有子模組的緩衝區。否則，僅產生作為此模組直接成員的緩衝區。預設值為 True。
remove_duplicate (bool, optional) – 是否移除結果中重複的緩衝區。預設值為 True。

產生：

(str, torch.Tensor) – 包含名稱和緩衝區的元組 (Tuple)

範例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, buf in self.named_buffers():
>>>     if name in ['running_var']:
>>>         print(buf.size())

named_children() → Iterator[Tuple[str, Module]]¶

回傳一個疊代器，遍歷直接子模組，產生模組的名稱和模組本身。

產生：: (str, Module) – 包含名稱和子模組的元組 (Tuple)

範例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules(memo: Optional[Set[Module]] = None, prefix: str = '', remove_duplicate: bool = True)¶

回傳一個疊代器，遍歷網路中的所有模組，產生模組的名稱和模組本身。

參數:

memo – 一個備忘錄 (memo)，用於儲存已新增到結果中的模組集合。
prefix – 將新增到模組名稱的前綴詞。
remove_duplicate – 是否移除結果中重複的模組實例。

產生：

(str, Module) – 名稱和模組的元組 (Tuple)

注意

重複的模組只會傳回一次。在以下範例中，l 只會傳回一次。

範例

>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters(prefix: str = '', recurse: bool = True, remove_duplicate: bool = True) → Iterator[Tuple[str, Parameter]]¶

回傳一個疊代器，遍歷模組參數 (parameter)，產生參數的名稱和參數本身。

參數:

prefix (str) – 要加在所有參數名稱前面的前綴詞。
recurse (bool) – 如果為 True，則產生此模組和所有子模組的參數。否則，僅產生作為此模組直接成員的參數。
remove_duplicate (bool, optional) – 是否移除結果中重複的參數。預設值為 True。

產生：

(str, Parameter) – 包含名稱和參數的元組 (Tuple)

範例

>>> # xdoctest: +SKIP("undefined vars")
>>> for name, param in self.named_parameters():
>>>     if name in ['bias']:
>>>         print(param.size())

property names¶

tensordict 的維度名稱。

名稱可以在建構時使用 names 參數進行設定。

另請參閱 refine_names() 以取得在建構後如何設定名稱的詳細資訊。

nanmean(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor¶

傳回輸入 tensordict 中所有非 NaN 元素的平均值。

參數:

dim (int, int 的 tuple, optional) – 如果 None，則傳回包含所有 leaves 平均值 (如果可以計算) 的無維度 tensordict。如果是整數或整數的 tuple，則只有當此維度與 tensordict 形狀相容時，才會在指定的維度上呼叫 mean。
keepdim (bool) – 輸出 tensor 是否保留 dim。

關鍵字參數:

dtype (torch.dtype, optional) – 傳回的 tensor 所需的資料類型。如果指定，則在執行操作之前，輸入 tensor 會被轉換為 dtype。這對於防止資料類型溢位很有用。預設值： None。
reduce (bool, optional) – 如果 True，則 reduciton 將發生在所有 TensorDict 值上，並傳回單個簡化的 tensor。預設為 False。

nansum(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor¶

傳回輸入 tensordict 中所有非 NaN 元素的總和。

參數:

dim (int, int 的 tuple, optional) – 如果為 None，則傳回一個無維度的 tensordict，其中包含所有葉節點的總和值 (如果可以計算)。如果是整數或整數元組，則僅當此維度與 tensordict 形狀相容時，才對指定的維度調用 sum 。
keepdim (bool) – 輸出 tensor 是否保留 dim。

關鍵字參數:

dtype (torch.dtype, optional) – 傳回的 tensor 所需的資料類型。如果指定，則在執行操作之前，輸入 tensor 會被轉換為 dtype。這對於防止資料類型溢位很有用。預設值： None。
reduce (bool, optional) – 如果 True，則 reduciton 將發生在所有 TensorDict 值上，並傳回單個簡化的 tensor。預設為 False。

property ndim: int¶: 請參閱 batch_dims()。

ndimension() → int¶: 請參閱 batch_dims()。

neg() → T¶: 計算 TensorDict 中每個元素的 neg() 值。

neg_() → T¶: 就地 (in-place) 計算 TensorDict 中每個元素的 neg() 值。

new_empty(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶

返回一個大小為 size 且包含空 tensors 的 TensorDict。

預設情況下，返回的 TensorDict 具有與此 TensorDict 相同的 torch.dtype 和 torch.device。

參數:

size (int...) – 一個整數的列表、元組或 torch.Size，用於定義輸出 tensor 的形狀。

關鍵字參數:

dtype (torch.dtype, optional) – 返回的 TensorDict 的所需類型。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 返回的 TensorDict 的所需設備。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄對返回的 tensors 的操作。預設值：False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需佈局。預設值：torch.strided。
pin_memory (bool, optional) – 如果設定，返回的 tensor 將分配在釘選的記憶體中。僅適用於 CPU tensors。預設值：False。

new_full(size: Size, fill_value, *, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶

傳回一個大小為 size 且以 1 填充的 TensorDict。

預設情況下，返回的 TensorDict 具有與此 TensorDict 相同的 torch.dtype 和 torch.device。

參數:

size (整數序列) – 一個整數的列表、元組或 torch.Size，用於定義輸出張量的形狀。
fill_value (純量) – 用於填充輸出張量的數值。

關鍵字參數:

dtype (torch.dtype, optional) – 返回的 TensorDict 的所需類型。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 返回的 TensorDict 的所需設備。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄對返回的 tensors 的操作。預設值：False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需佈局。預設值：torch.strided。
pin_memory (bool, optional) – 如果設定，返回的 tensor 將分配在釘選的記憶體中。僅適用於 CPU tensors。預設值：False。

new_ones(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶

傳回一個大小為 size 且以 1 填充的 TensorDict。

預設情況下，返回的 TensorDict 具有與此 TensorDict 相同的 torch.dtype 和 torch.device。

參數:

size (int...) – 一個整數的列表、元組或 torch.Size，用於定義輸出 tensor 的形狀。

關鍵字參數:

dtype (torch.dtype, optional) – 返回的 TensorDict 的所需類型。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 返回的 TensorDict 的所需設備。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄對返回的 tensors 的操作。預設值：False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需佈局。預設值：torch.strided。
pin_memory (bool, optional) – 如果設定，返回的 tensor 將分配在釘選的記憶體中。僅適用於 CPU tensors。預設值：False。

new_tensor(data: torch.Tensor | tensordict.base.TensorDictBase, *, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, pin_memory: Optional[bool] = None)¶

返回一個新的 TensorDict，其資料為 tensor data。

預設情況下，返回的 TensorDict 值與此 tensor 具有相同的 torch.dtype 和 torch.device。

data 也可以是一個 tensor 集合 (TensorDict 或 tensorclass)，在這種情況下，new_tensor 方法會迭代 self 和 data 的 tensor 對。

參數:

data (torch.Tensor 或 TensorDictBase) – 要複製的資料。

關鍵字參數:

dtype (torch.dtype, optional) – 返回的 TensorDict 的所需類型。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 返回的 TensorDict 的所需設備。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄對返回的 tensors 的操作。預設值：False。
pin_memory (bool, optional) – 如果設定，返回的 tensor 將分配在釘選的記憶體中。僅適用於 CPU tensors。預設值：False。

new_zeros(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶

返回一個大小為 size 且填充 0 的 TensorDict。

預設情況下，返回的 TensorDict 具有與此 TensorDict 相同的 torch.dtype 和 torch.device。

參數:

size (int...) – 一個整數的列表、元組或 torch.Size，用於定義輸出 tensor 的形狀。

關鍵字參數:

dtype (torch.dtype, optional) – 返回的 TensorDict 的所需類型。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 返回的 TensorDict 的所需設備。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄對返回的 tensors 的操作。預設值：False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需佈局。預設值：torch.strided。
pin_memory (bool, optional) – 如果設定，返回的 tensor 將分配在釘選的記憶體中。僅適用於 CPU tensors。預設值：False。

non_tensor_items(include_nested: bool = False)¶: 傳回所有非 tensor 的葉節點，可能以遞迴方式進行。

norm(*, out=None, dtype: torch.dtype | None = None)¶

計算 tensordict 中每個 tensor 的範數。

關鍵字參數:

out (TensorDict, optional) – 輸出 tensordict。
dtype (torch.dtype, optional) – 輸出 dtype (torch>=2.4)。

numel() → int¶

批次中的元素總數。

下限為 1，因為具有空形狀的兩個 tensordict 的堆疊將有兩個元素，因此我們認為 tensordict 至少有 1 個元素大。

numpy()¶

將 tensordict 轉換為 (可能為巢狀) numpy 陣列的字典。

非 tensor 資料會以原樣暴露。

範例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({"a": {"b": torch.zeros(()), "c": "a string!"}})
>>> print(data)
TensorDict(
    fields={
        a: TensorDict(
            fields={
                b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                c: NonTensorData(data=a string!, batch_size=torch.Size([]), device=None)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(data.numpy())
{'a': {'b': array(0., dtype=float32), 'c': 'a string!'}}

param_count(*, count_duplicates: bool = True) → int¶

計算參數的數量（可索引項目的總數），僅考慮 tensor。

關鍵字參數:: count_duplicates (bool) – 是否將重複的 tensor 計算為獨立的 tensor。如果 False，則僅會丟棄嚴格相同的 tensors（來自共同基礎 tensor 的相同檢視但具有不同的 ids 將被計算兩次）。預設為 True（假設每個 tensor 都是單一副本）。

parameters(recurse: bool = True) → Iterator[Parameter]¶

傳回模組參數的迭代器。

這通常傳遞給優化器。

參數:: recurse (bool) – 如果為 True，則產生此模組和所有子模組的參數。否則，僅產生作為此模組直接成員的參數。
產生：: Parameter – 模組參數

範例

>>> # xdoctest: +SKIP("undefined vars")
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

permute(*args, **kwargs)¶

傳回 tensordict 的視圖，其批次維度根據 dims 重新排列。

參數:

*dims_list (int) – tensordict 批次維度的新順序。或者，可以提供單個整數的可迭代物件。
dims (list of int) – 呼叫 permute(…) 的替代方法。

回傳值:

一個新的 tensordict，其批次維度具有所需的順序。

範例

>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4])
>>> print(tensordict.permute([1, 0]))
PermutedTensorDict(
    source=TensorDict(
        fields={
            a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)},
        batch_size=torch.Size([3, 4]),
        device=cpu,
        is_shared=False),
    op=permute(dims=[1, 0]))
>>> print(tensordict.permute(1, 0))
PermutedTensorDict(
    source=TensorDict(
        fields={
            a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)},
        batch_size=torch.Size([3, 4]),
        device=cpu,
        is_shared=False),
    op=permute(dims=[1, 0]))
>>> print(tensordict.permute(dims=[1, 0]))
PermutedTensorDict(
    source=TensorDict(
        fields={
            a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)},
        batch_size=torch.Size([3, 4]),
        device=cpu,
        is_shared=False),
    op=permute(dims=[1, 0]))

pin_memory(*args, **kwargs)¶

在儲存的 tensor 上呼叫 pin_memory()。

參數:

num_threads (int or str) – 如果提供，則用於在葉子上呼叫 pin_memory 的執行緒數。預設為 None，這會在 ThreadPoolExecutor(max_workers=None) 中設定大量的執行緒。若要在主執行緒上執行對 pin_memory() 的所有呼叫，請傳遞 num_threads=0。
inplace (bool, optional) – 如果 True，則 tensordict 會就地修改。預設為 False。

pin_memory_(num_threads: int | str = 0) → T¶

在儲存的 tensor 上呼叫 pin_memory() 並傳回就地修改的 TensorDict。

參數:: num_threads (int or str) – 如果提供，則用於在葉子上呼叫 pin_memory 的執行緒數。如果傳遞 "auto"，則會自動判斷執行緒數。

pop(key: NestedKey, default: Any = _NoDefault.ZERO) → Tensor¶

從 tensordict 中移除並傳回一個值。

如果該值不存在且沒有提供預設值，則會拋出 KeyError。

參數:

key (str 或 巢狀鍵) – 要尋找的條目。
default (Any, 可選) – 如果找不到鍵，則傳回的值。

範例

>>> td = TensorDict({"1": 1}, [])
>>> one = td.pop("1")
>>> assert one == 1
>>> none = td.pop("1", default=None)
>>> assert none is None

popitem()¶

移除最後插入 TensorDict 的項目。

popitem 只會傳回非巢狀的值。

pow(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T¶

計算 self 中每個元素與 other 的次方，並傳回包含結果的 tensor。

other 可以是單一的 float 數字、Tensor 或 TensorDict。

當 other 是一個 tensor 時，input 和 other 的形狀必須是可以廣播的。

參數:: other (float, tensor 或 tensordict) – 指數值
關鍵字參數:: default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

pow_(other: tensordict.base.TensorDictBase | torch.Tensor) → T¶: pow() 的原地 (in-place) 版本。

注意

Inplace pow 不支援 default 關鍵字引數。

prod(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor¶

傳回輸入 tensordict 中所有元素的數值乘積。

參數:

dim (int, int 的 tuple, optional) – 如果 None，則傳回包含所有 leaf 乘積值的無維度 tensordict（如果可以計算）。如果是整數或整數的 tuple，則只有在該維度與 tensordict 形狀相容時，才對指定的維度呼叫 prod。
keepdim (bool) – 輸出 tensor 是否保留 dim。

關鍵字參數:

dtype (torch.dtype, optional) – 傳回的 tensor 所需的資料類型。如果指定，則在執行操作之前，輸入 tensor 會被轉換為 dtype。這對於防止資料類型溢位很有用。預設值： None。
reduce (bool, optional) – 如果 True，則 reduciton 將發生在所有 TensorDict 值上，並傳回單個簡化的 tensor。預設為 False。

qint32()¶: 將所有 tensors 轉換為 torch.qint32。

qint8()¶: 將所有 tensors 轉換為 torch.qint8。

quint4x2()¶: 將所有 tensors 轉換為 torch.quint4x2。

quint8()¶: 將所有 tensors 轉換為 torch.quint8。

reciprocal() → T¶: 計算 TensorDict 中每個元素的 reciprocal() 值。

reciprocal_() → T¶: 就地計算 TensorDict 中每個元素的 reciprocal() 值。

record_stream(stream: Stream)¶

將 tensordict 標記為已由此串流使用。

當 tensordict 被釋放時，請確保在釋放時排隊在串流上的所有工作完成之前，張量記憶體不會被重複用於其他張量。

請參閱 record_stream() 以取得更多資訊。

recv(src: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → int¶

接收 tensordict 的內容，並使用該內容更新內容。

查看 send 方法中的範例以取得內容。

參數:

src (int) – 來源工作者的 rank。

關鍵字參數:

group (torch.distributed.ProcessGroup, 選用) – 如果設定，指定的進程群組將用於通訊。否則，將使用預設的進程群組。預設為 None。
init_tag (int) – 來源工作者使用的 init_tag。
pseudo_rand (bool) – 若為 True，則標籤序列將為偽隨機，允許從不同節點傳送多個資料而不會重疊。請注意，產生這些偽隨機數字的成本很高 (1e-5 秒/數字)，這意味著它可能會降低演算法的執行時間。此值必須與傳遞給 send() 的值相符。預設值為 False。

reduce(dst, op=None, async_op=False, return_premature=False, group=None)¶

在所有機器上簡化 tensordict。

只有 rank 為 dst 的程序才會收到最終結果。

refine_names(*names) → T¶

根據 names 精煉 self 的維度名稱。

精煉是重命名的特殊情況，它會「提升」未命名的維度。 None 維度可以精煉為具有任何名稱；命名的維度只能精煉為具有相同的名稱。

由於命名的張量可以與未命名的張量共存，因此精煉名稱提供了一種很好的方法來編寫可以使用命名和未命名張量的、感知命名張量的程式碼。

names 最多可能包含一個 Ellipsis (…)。 Ellipsis 會貪婪地展開；它會就地展開，以使用來自 self.names 相應索引的名稱，使 names 的長度與 self.dim() 相同。

傳回值：與輸入名稱對應的具名維度的相同 tensordict。

範例

>>> td = TensorDict({}, batch_size=[3, 4, 5, 6])
>>> tdr = td.refine_names(None, None, None, "d")
>>> assert tdr.names == [None, None, None, "d"]
>>> tdr = td.refine_names("a", None, None, "d")
>>> assert tdr.names == ["a", None, None, "d"]

register_backward_hook(hook: Callable[[Module, Union[Tuple[Tensor, ...], Tensor], Union[Tuple[Tensor, ...], Tensor]], Union[None, Tuple[Tensor, ...], Tensor]]) → RemovableHandle¶

在模組上註冊一個 backward hook。

此函式已被棄用，建議使用 register_full_backward_hook()，並且此函式的行為將在未來的版本中變更。

回傳值:: 一個句柄，可用於通過調用 handle.remove() 來刪除新增的 hook。
回傳類型:: torch.utils.hooks.RemovableHandle

register_buffer(name: str, tensor: Optional[Tensor], persistent: bool = True) → None¶

將 buffer 新增到模組中。

這通常用於註冊不應被視為模型參數的 buffer。例如，BatchNorm 的 running_mean 不是參數，而是模組狀態的一部分。預設情況下，Buffers 是持久性的，並將與參數一起保存。可以透過將 persistent 設置為 False 來更改此行為。持久性 buffer 和非持久性 buffer 之間的唯一區別是後者不會成為此模組的 state_dict 的一部分。

可以使用給定的名稱作為屬性來存取 Buffers。

參數:

name (str) – buffer 的名稱。可以使用給定的名稱從此模組存取 buffer
tensor (Tensor 或 None) – 要註冊的 buffer。如果 None，則忽略在 buffers 上運行的操作，例如 cuda。如果 None，則該 buffer 不包含在模組的 state_dict 中。
persistent (bool) – buffer 是否為此模組的 state_dict 的一部分。

範例

>>> # xdoctest: +SKIP("undefined vars")
>>> self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook(hook: Union[Callable[[T, Tuple[Any, ...], Any], Optional[Any]], Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]]], *, prepend: bool = False, with_kwargs: bool = False, always_call: bool = False) → RemovableHandle¶

在模組上註冊一個 forward hook。

這個 hook 會在每次 forward() 計算出輸出後被呼叫。

如果 with_kwargs 是 False 或未指定，則輸入僅包含傳遞給模組的位置參數。關鍵字參數不會傳遞給 hooks，只會傳遞給 forward。這個 hook 可以修改輸出。它可以就地修改輸入，但由於它是在 forward() 被呼叫後才被呼叫，因此不會對 forward 產生影響。這個 hook 應該具有以下簽章：

hook(module, args, output) -> None or modified output

如果 with_kwargs 是 True，forward hook 將被傳遞給 forward 函數的 kwargs，並且預期會回傳可能經過修改的輸出。這個 hook 應該具有以下簽章：

hook(module, args, kwargs, output) -> None or modified output

參數:

hook (Callable) – 要註冊的使用者定義 hook。
prepend (bool) – 如果 True，則提供的 hook 將在這個 torch.nn.modules.Module 上所有現有的 forward hooks 之前觸發。否則，提供的 hook 將在這個 torch.nn.modules.Module 上所有現有的 forward hooks 之後觸發。請注意，使用 register_module_forward_hook() 註冊的全域 forward hooks 將在透過此方法註冊的所有 hooks 之前觸發。預設值：False
with_kwargs (bool) – 如果 True，則 hook 將被傳遞給 forward 函數的 kwargs。預設值：False
always_call (bool) – 如果 True，則無論在呼叫 Module 時是否引發例外，都會執行 hook。預設值：False

回傳值:

一個句柄，可用於通過調用 handle.remove() 來刪除新增的 hook。

回傳類型:

torch.utils.hooks.RemovableHandle

register_forward_pre_hook(hook: Union[Callable[[T, Tuple[Any, ...]], Optional[Any]], Callable[[T, Tuple[Any, ...], Dict[str, Any]], Optional[Tuple[Any, Dict[str, Any]]]]], *, prepend: bool = False, with_kwargs: bool = False) → RemovableHandle¶

在模組上註冊一個 forward pre-hook。

這個 hook 會在每次呼叫 forward() 之前被呼叫。

如果 with_kwargs 為 false 或未指定，則輸入僅包含傳遞給模組的位置引數。關鍵字引數將不會傳遞給 hooks，而只會傳遞給 forward。hook 可以修改輸入。使用者可以在 hook 中回傳一個 tuple 或單一修改後的值。如果回傳的是單一值（除非該值已經是一個 tuple），我們將會把該值包裝成一個 tuple。hook 應該具有以下簽章：

hook(module, args) -> None or modified input

如果 with_kwargs 為 true，則 forward pre-hook 將會被傳遞給 forward 函數的 kwargs。並且如果 hook 修改了輸入，則 args 和 kwargs 都應該被回傳。hook 應該具有以下簽章：

hook(module, args, kwargs) -> None or a tuple of modified input and kwargs

參數:

hook (Callable) – 要註冊的使用者定義 hook。
prepend (bool) – 如果為 true，則提供的 hook 將會在該 torch.nn.modules.Module 上所有現有的 forward_pre hooks 之前觸發。否則，提供的 hook 將會在該 torch.nn.modules.Module 上所有現有的 forward_pre hooks 之後觸發。請注意，透過 register_module_forward_pre_hook() 註冊的全域 forward_pre hooks 將會在所有透過此方法註冊的 hooks 之前觸發。預設值：False
with_kwargs (bool) – 如果為 true，則 hook 將會被傳遞給 forward 函數的 kwargs。預設值：False

回傳值:

一個句柄，可用於通過調用 handle.remove() 來刪除新增的 hook。

回傳類型:

torch.utils.hooks.RemovableHandle

register_full_backward_hook(hook: Callable[[Module, Union[Tuple[Tensor, ...], Tensor], Union[Tuple[Tensor, ...], Tensor]], Union[None, Tuple[Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle¶

在模組上註冊一個 backward hook。

每當計算關於模組的梯度時，就會呼叫這個 hook，也就是說，只有在計算關於模組輸出的梯度時，hook 才會執行。hook 應該具有以下簽章：

hook(module, grad_input, grad_output) -> tuple(Tensor) or None

grad_input 和 grad_output 是包含關於輸入和輸出的梯度的 tuples。hook 不應該修改其引數，但它可以選擇性地回傳關於輸入的新梯度，該梯度將會被用來代替後續計算中的 grad_input。grad_input 只會對應到作為位置引數給定的輸入，並且所有 kwarg 引數都會被忽略。grad_input 和 grad_output 中的條目對於所有非 Tensor 引數將會是 None。

由於技術原因，當此 hook 應用於 Module 時，其 forward 函數將會收到傳遞給 Module 的每個 Tensor 的視圖 (view)。同樣地，呼叫者將會收到由 Module 的 forward 函數回傳的每個 Tensor 的視圖 (view)。

警告

使用 backward hooks 時，不允許就地修改輸入或輸出，否則會引發錯誤。

參數:

hook (Callable) – 要註冊的使用者定義 hook。
prepend ( bool) – 若為 true，則提供的 hook 將會在該 torch.nn.modules.Module 上的所有現有 backward hook 之前觸發。否則，提供的 hook 將會在該 torch.nn.modules.Module 上的所有現有 backward hook 之後觸發。請注意，使用 register_module_full_backward_hook() 註冊的全域 backward hook 將會在透過此方法註冊的所有 hook 之前觸發。

回傳值:

一個句柄，可用於通過調用 handle.remove() 來刪除新增的 hook。

回傳類型:

torch.utils.hooks.RemovableHandle

register_full_backward_pre_hook(hook: Callable[[Module, Union[Tuple[Tensor, ...], Tensor]], Union[None, Tuple[Tensor, ...], Tensor]], prepend: bool = False) → RemovableHandle¶

在模組上註冊一個 backward pre-hook。

每次計算模組的梯度時，都會呼叫這個 hook。這個 hook 應該具有以下簽名：

hook(module, grad_output) -> tuple[Tensor] or None

grad_output 是一個 tuple。這個 hook 不應該修改它的參數，但它可以選擇性地返回一個新的輸出梯度，用來代替後續計算中的 grad_output。grad_output 中的所有非 Tensor 參數，其值將會是 None。

由於技術原因，當此 hook 應用於 Module 時，其 forward 函數將會收到傳遞給 Module 的每個 Tensor 的視圖 (view)。同樣地，呼叫者將會收到由 Module 的 forward 函數回傳的每個 Tensor 的視圖 (view)。

警告

當使用 backward hook 時，不允許就地修改輸入，否則會引發錯誤。

參數:

hook (Callable) – 要註冊的使用者定義 hook。
prepend ( bool) – 若為 true，則提供的 hook 將會在該 torch.nn.modules.Module 上的所有現有 backward_pre hook 之前觸發。否則，提供的 hook 將會在該 torch.nn.modules.Module 上的所有現有 backward_pre hook 之後觸發。請注意，使用 register_module_full_backward_pre_hook() 註冊的全域 backward_pre hook 將會在透過此方法註冊的所有 hook 之前觸發。

回傳值:

一個句柄，可用於通過調用 handle.remove() 來刪除新增的 hook。

回傳類型:

torch.utils.hooks.RemovableHandle

register_get_post_hook(hook)¶: 註冊一個 hook，以便在對 leaf tensor 執行任何 get 操作後呼叫。

register_load_state_dict_post_hook(hook)¶

註冊一個 post-hook，以便在模組的 load_state_dict() 被呼叫後執行。

它應該具有以下簽名：: hook(module, incompatible_keys) -> None

module 參數是註冊此 hook 的目前模組，而 incompatible_keys 參數是一個 NamedTuple，包含屬性 missing_keys 和 unexpected_keys。missing_keys 是一個 list，其中包含遺失的 key 的 str，而 unexpected_keys 是一個 list，其中包含意外 key 的 str。

如果需要，可以就地修改給定的 incompatible_keys。

請注意，當使用 strict=True 呼叫 load_state_dict() 時執行的檢查，會受到 hook 對 missing_keys 或 unexpected_keys 所做的修改影響，如同預期。新增至任一組鍵將導致在 strict=True 時拋出錯誤，而清除遺失和未預期的鍵將避免錯誤。

回傳值:: 一個句柄，可用於通過調用 handle.remove() 來刪除新增的 hook。
回傳類型:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)¶

註冊一個 pre-hook，使其在模組的 load_state_dict() 被呼叫之前執行。

它應該具有以下簽名：: hook(module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) -> None # noqa: B950

參數:: hook (Callable) – 可呼叫的 hook，將在載入 state dict 之前被調用。

register_module(name: str, module: Optional[Module]) → None¶: 是 add_module() 的別名。

register_parameter(name: str, param: Optional[Parameter]) → None¶

新增一個參數到模組中。

可以使用給定的名稱作為屬性來存取該參數。

參數:

name (str) – 參數的名稱。可以使用給定的名稱從此模組存取該參數。
param (Parameter 或 None) – 要新增到模組的參數。如果 None，則忽略在參數上執行的操作，例如 cuda。如果 None，則參數不會包含在模組的 state_dict 中。

register_state_dict_post_hook(hook)¶

註冊 state_dict() 方法的 post-hook。

它應該具有以下簽名：: hook(module, state_dict, prefix, local_metadata) -> None

註冊的 hook 可以就地修改 state_dict。

register_state_dict_pre_hook(hook)¶

註冊 state_dict() 方法的 pre-hook。

它應該具有以下簽名：: hook(module, prefix, keep_vars) -> None

註冊的 hook 可用於在呼叫 state_dict 之前執行預處理。

rename(*names, **rename_map)¶

返回一個重新命名維度的 tensordict 的副本。

範例

>>> td = TensorDict({}, batch_size=[1, 2, 3 ,4])
>>> td.names = list("abcd")
>>> td_rename = td.rename(c="g")
>>> assert td_rename.names == list("abgd")

rename_(*names, **rename_map)¶

與 rename() 相同，但會就地執行重新命名。

範例

>>> td = TensorDict({}, batch_size=[1, 2, 3 ,4])
>>> td.names = list("abcd")
>>> assert td.rename_(c="g")
>>> assert td.names == list("abgd")

rename_key_(old_key: NestedKey, new_key: NestedKey, safe: bool = False) → TensorDictBase¶

使用新的字串重新命名鍵，並傳回具有更新鍵名稱的相同 tensordict。

參數:

old_key (str 或 nested key) – 要重新命名的鍵。
new_key (str 或 nested key) – 條目的新名稱。
safe (bool, optional) – 如果 True，則當新的鍵已存在於 TensorDict 中時，會拋出錯誤。

回傳值:

self

replace(*args, **kwargs)¶

建立 tensordict 的淺層副本，其中的條目已被替換。

接受一個未命名的參數，該參數必須是 TensorDictBase 子類的字典。此外，可以使用命名的關鍵字參數更新第一層條目。

回傳值:: 如果輸入為非空，則為 self 的副本，其中包含更新的條目。如果提供了空字典或未提供字典，並且 kwargs 為空，則傳回 self。

requires_grad_(requires_grad=True) → T¶

更改是否自動微分 (autograd) 應記錄在此張量上的操作：就地設定此張量的 requires_grad 屬性。

傳回此 tensordict。

參數:: requires_grad (bool, optional) – 自動微分是否應記錄在此 tensordict 上的操作。預設為 True。

reshape(*shape: int)¶

傳回具有所需形狀的連續、重新塑形的張量。

參數:: *shape (int) – 結果 tensordict 的新形狀。
回傳值:: 具有重新塑形鍵的 TensorDict

範例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td = td.reshape(12)
>>> print(td['x'])
torch.Tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])

round() → T¶: 計算 TensorDict 每個元素的 round() 值。

round_() → T¶: 就地計算 TensorDict 每個元素的 round() 值。

save(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

將 tensordict 儲存到磁碟。

此函式是 memmap() 的代理。

property saved_path¶

傳回儲存 memmap saved TensorDict 的路徑。

一旦 is_memmap() 傳回 False (例如，當 tensordict 被解鎖時)，此引數即消失。

select(*keys: NestedKey, inplace: bool = False, strict: bool = True) → T¶

選擇 tensordict 的鍵，並傳回一個僅包含所選鍵的新 tensordict。

數值不會被複製：對原始或新的 tensordict 的張量進行原地 (in-place) 修改將會導致兩個 tensordict 都發生變化。

參數:

*keys (str) – 要選擇的鍵。
inplace (bool) – 如果為 True，則 tensordict 會原地修剪。預設值為 False。
strict (bool, optional) – 選擇不存在的鍵是否會傳回錯誤。預設值：True。

回傳值:

一個新的 tensordict (如果 inplace=True 則為同一個 tensordict)，僅包含所選鍵。

注意

若要在 tensordict 中選擇鍵，並傳回一個不包含這些鍵的版本，請參閱 split_keys() 方法。

範例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, [])
>>> td.select("a", ("b", "c"))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.select("a", "b")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.select("this key does not exist", strict=False)
TensorDict(
    fields={
    },
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

send(dst: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → None¶

將 tensordict 的內容傳送到遠端 worker。

參數:

dst (int) – 內容應發送到的目標 worker 的排名。

關鍵字參數:

group (torch.distributed.ProcessGroup, 選用) – 如果設定，指定的進程群組將用於通訊。否則，將使用預設的進程群組。預設為 None。
init_tag (int) – 用於標記張量的初始標籤。請注意，這將根據 TensorDict 中包含的張量數量進行遞增。
pseudo_rand (bool) – 如果為 True，標籤的序列將是偽隨機的，允許從不同的節點發送多個資料而不會重疊。請注意，這些偽隨機數的產生成本很高（1e-5 秒/數字），這意味著它可能會降低演算法的運行時間。預設值為 False。

範例

>>> from torch import multiprocessing as mp
>>> from tensordict import TensorDict
>>> import torch
>>>
>>>
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...
...     td = TensorDict(
...         {
...             ("a", "b"): torch.randn(2),
...             "c": torch.randn(2, 3),
...             "_": torch.ones(2, 1, 5),
...         },
...         [2],
...     )
...     td.send(0)
...
>>>
>>> def server(queue):
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=0,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...     td = TensorDict(
...         {
...             ("a", "b"): torch.zeros(2),
...             "c": torch.zeros(2, 3),
...             "_": torch.zeros(2, 1, 5),
...         },
...         [2],
...     )
...     td.recv(1)
...     assert (td != 0).all()
...     queue.put("yuppie")
...
>>>
>>> if __name__=="__main__":
...     queue = mp.Queue(1)
...     main_worker = mp.Process(target=server, args=(queue,))
...     secondary_worker = mp.Process(target=client)
...
...     main_worker.start()
...     secondary_worker.start()
...     out = queue.get(timeout=10)
...     assert out == "yuppie"
...     main_worker.join()
...     secondary_worker.join()

set(key: NestedKey, item: Tensor, inplace: bool = False, **kwargs: Any) → TensorDictBase¶

設定新的鍵值對。

參數:

key (str, str 的 tuple) – 要設定的鍵的名稱。
item (torch.Tensor 或 等效物件, TensorDictBase 實例) – 要儲存在 tensordict 中的值。
inplace (bool, optional) – 如果 True 且鍵與 tensordict 中現有的鍵相符，則該鍵值對的更新將會就地發生。如果 inplace 為 True 且找不到該條目，則會新增它。對於更嚴格的就地操作，請改用 set_()。預設值為 False。

關鍵字參數:

non_blocking (bool, optional) – 如果 True 且此複製發生在不同裝置之間，則複製可能會相對於主機非同步發生。

回傳值:

self

範例

>>> td = TensorDict({}, batch_size[3, 4])
>>> td.set("x", torch.randn(3, 4))
>>> y = torch.randn(3, 4, 5)
>>> td.set("y", y, inplace=True) # works, even if 'y' is not present yet
>>> td.set("y", torch.zeros_like(y), inplace=True)
>>> assert (y==0).all() # y values are overwritten
>>> td.set("y", torch.ones(5), inplace=True) # raises an exception as shapes mismatch

set_(key: NestedKey, item: Tensor) → T¶

設定一個值到已存在的鍵，同時保留原始儲存空間。

參數:

key (str) – 值的名稱
item (torch.Tensor 或 相容類型, TensorDictBase) – 要儲存在 tensordict 中的值

關鍵字參數:

non_blocking (bool, optional) – 如果 True 且此複製發生在不同裝置之間，則複製可能會相對於主機非同步發生。

回傳值:

self

範例

>>> td = TensorDict({}, batch_size[3, 4])
>>> x = torch.randn(3, 4)
>>> td.set("x", x)
>>> td.set_("x", torch.zeros_like(x))
>>> assert (x == 0).all()

set_at_(key: NestedKey, value: Tensor, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]]) → T¶

在 index 指示的索引位置就地設定值。

參數:

key (str, str 的 tuple) – 要修改的鍵。
value (torch.Tensor) – 要在索引 index 設定的值
index (int, tensor 或 tuple) – 寫入值的索引。

關鍵字參數:

non_blocking (bool, optional) – 如果 True 且此複製發生在不同裝置之間，則複製可能會相對於主機非同步發生。

回傳值:

self

範例

>>> td = TensorDict({}, batch_size[3, 4])
>>> x = torch.randn(3, 4)
>>> td.set("x", x)
>>> td.set_at_("x", value=torch.ones(1, 4), index=slice(1))
>>> assert (x[0] == 1).all()

set_extra_state(state: Any) → None¶

設定載入的 state_dict 中包含的額外狀態。

此函數從 load_state_dict() 呼叫，以處理在 state_dict 中找到的任何額外狀態。如果您需要在模組的 state_dict 中儲存額外狀態，請實現此函數和相應的 get_extra_state()。

參數:: state (dict) – 來自 state_dict 的額外狀態

set_non_tensor(key: NestedKey, value: Any)¶

使用 tensordict.tensorclass.NonTensorData 在 tensordict 中註冊一個非 tensor 值。

可以使用 TensorDictBase.get_non_tensor() 或直接使用 get 檢索該值，它將返回 tensordict.tensorclass.NonTensorData 物件。

return: self

範例

>>> data = TensorDict({}, batch_size=[])
>>> data.set_non_tensor(("nested", "the string"), "a string!")
>>> assert data.get_non_tensor(("nested", "the string")) == "a string!"
>>> # regular `get` works but returns a NonTensorData object
>>> data.get(("nested", "the string"))
NonTensorData(
    data='a string!',
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

set_submodule(target: str, module: Module) → None¶

如果由 target 給定的子模組存在，則設定它，否則拋出錯誤。

例如，假設您有一個 nn.Module A，如下所示

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

（該圖顯示了一個 nn.Module A。A 有一個巢狀的子模組 net_b，它本身有兩個子模組 net_c 和 linear。net_c 然後有一個子模組 conv。）

若要使用新的子模組 Linear 覆寫 Conv2d，您可以呼叫 set_submodule("net_b.net_c.conv", nn.Linear(33, 16))。

參數:

target – 要查找的子模組的完整字串名稱。（有關如何指定完整字串，請參閱上面的範例。）
module – 要設定子模組的模組。

引發:

ValueError – 如果目標字串為空。
AttributeError – 如果目標字串引用無效的路徑，或解析為不是 nn.Module 的東西

setdefault(key: NestedKey, default: Tensor, inplace: bool = False) → Tensor¶

如果 key 不在 tensordict 中，則插入 key 條目，其值為 default。

如果 key 在 tensordict 中，則傳回 key 的值，否則傳回 default。

參數:

key (str 或 巢狀鍵 (nested key)) – 值的名稱。
default (torch.Tensor 或 相容類型 (compatible type), TensorDictBase) – 如果 tensordict 中尚未存在該鍵，則要儲存在 tensordict 中的值。

回傳值:

tensordict 中 key 的值。如果先前未設定該鍵，則將為 default。

範例

>>> td = TensorDict({}, batch_size=[3, 4])
>>> val = td.setdefault("a", torch.zeros(3, 4))
>>> assert (val == 0).all()
>>> val = td.setdefault("a", torch.ones(3, 4))
>>> assert (val == 0).all() # output is still 0

property shape: Size¶: 請參閱 batch_size。

share_memory() → T¶: 請參閱 torch.Tensor.share_memory_()。

share_memory_(*args, **kwargs)¶

將所有 tensors 放置在共享記憶體中。

然後會鎖定 TensorDict，這表示任何非原地 (in-place) 的寫入操作都會拋出例外（例如，重新命名、設定或移除條目）。相反地，一旦 tensordict 被解鎖，share_memory 屬性就會變成 False，因為不再保證跨程序的相同性 (cross-process identity)。

回傳值:: self

sigmoid() → T¶: 計算 TensorDict 中每個元素的 sigmoid() 值。

sigmoid_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 sigmoid() 值。

sign() → T¶: 計算 TensorDict 中每個元素的 sign() 值。

sign_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 sign() 值。

sin() → T¶: 計算 TensorDict 中每個元素的 sin() 值。

sin_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 sin() 值。

sinh() → T¶: 計算 TensorDict 中每個元素的 sinh() 值。

sinh_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 sinh() 值。

size(dim: Optional[int] = None) → torch.Size | int¶

回傳 dim 所指示的維度的大小。

如果未指定 dim，則回傳 TensorDict 的 batch_size 屬性。

property sorted_keys: list[tensordict._nestedkey.NestedKey]¶

回傳依字母順序排序的鍵。

不支援額外的引數。

如果 TensorDict 已鎖定，則鍵會被快取，直到 TensorDict 解鎖以加快執行速度。

split(split_size: int | list[int], dim: int = 0) → list[tensordict.base.TensorDictBase]¶

使用指定的大小在給定的維度中分割 TensorDict 中的每個張量，類似於 torch.split。

回傳一個 TensorDict 實例的列表，其中包含分割區塊的項目視圖。

參數:

split_size (int 或 List(int)) – 單個區塊的大小或每個區塊的大小列表。
dim (int) – 沿其分割張量的維度。

回傳值:

具有指定大小在給定維度中的 TensorDict 列表。

範例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td0, td1 = td.split([1, 2], dim=0)
>>> print(td0['x'])
torch.Tensor([[0, 1, 2, 3]])

split_keys(*key_sets, inplace=False, strict: bool = True, reproduce_struct: bool = False)¶

將 tensordict 分割成子集，給定一組或多組鍵。

此方法將回傳 N+1 個 tensordict，其中 N 是提供的引數數量。

參數:

inplace (bool, optional) – 如果 True，則鍵會從 self 原地 (in-place) 移除。預設值為 False。
strict (bool, optional) – 如果 True，則在遺失鍵時會引發例外。預設值為 True。
reproduce_struct (bool, optional) – 如果 True，則即使某些子 tensordict 不包含任何葉節點，所有回傳的 tensordict 都具有與 self 相同的樹狀結構。

注意

None 非張量值將被忽略且不會回傳。

注意

此方法不會檢查提供的列表中的重複項。

範例

>>> td = TensorDict(
...     a=0,
...     b=0,
...     c=0,
...     d=0,
... )
>>> td_a, td_bc, td_d = td.split_keys(["a"], ["b", "c"])
>>> print(td_bc)

sqrt()¶: 計算 self 的元素平方根。

sqrt_()¶: sqrt() 的原地 (in-place) 版本。

squeeze(*args, **kwargs)¶

壓縮介於 -self.batch_dims+1 和 self.batch_dims-1 之間的所有維度的 tensor，並在新的 tensordict 中傳回它們。

參數:: dim (Optional[int]) – 要沿其壓縮的維度。如果 dim 是 None，則將壓縮所有單例維度。預設值為 None。

範例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> td = td.squeeze()
>>> td.shape
torch.Size([3, 4])
>>> td.get("x").shape
torch.Size([3, 4, 2])

此操作也可以用作上下文管理器。對原始 tensordict 的變更將會異地 (out-place) 發生，即原始 tensor 的內容不會被變更。這也假設 tensordict 未被鎖定（否則，需要先解鎖 tensordict）。此功能與隱式壓縮不相容。

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> with td.squeeze(1) as tds:
...     tds.set("y", torch.zeros(3, 4))
>>> assert td.get("y").shape == [3, 1, 4]

classmethod stack(input, dim=0, *, out=None)¶

沿給定維度將 tensordict 堆疊到單個 tensordict 中。

此呼叫等同於呼叫 torch.stack()，但與 torch.compile 相容。

stack_from_tensordict(dim: int = 0, *, sorted: Optional[Union[bool, List[NestedKey]]] = None, out: Optional[Tensor] = None) → Tensor¶

將 tensordict 的所有條目堆疊到單個 tensor 中。

參數:

dim (int, optional) – 條目應沿其堆疊的維度。

關鍵字參數:

sorted (bool 或 NestedKeys 列表) – 如果 True，條目將按字母順序堆疊。如果 False （預設），將使用 dict 順序。或者，可以提供金鑰名稱的清單，並且 tensor 將相應地堆疊。這會產生一些額外負擔，因為將根據 tensordict 中的 leaf 名稱清單檢查金鑰清單。
out (torch.Tensor, optional) – 用於堆疊操作的可選目的地 tensor。

stack_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → T¶

將條目堆疊到新條目中，並且可能會移除原始值。

參數:: keys (NestedKey 序列) – 要堆疊的條目。

關鍵字參數

out_key (NestedKey): 堆疊輸入的新金鑰名稱。 keep_entries (bool, optional): 如果 False，則會刪除 keys 中的條目。

預設為 False。

dim (int, optional): 堆疊必須沿其發生的維度。: 預設為 0。

Returns: self

範例

>>> td = TensorDict(a=torch.zeros(()), b=torch.ones(()))
>>> td.stack_tensors("a", "b", out_key="c")
>>> assert "a" not in td
>>> assert (td["c"] == torch.tensor([0, 1])).all()

state_dict(*args, destination=None, prefix='', keep_vars=False, flatten=True)¶

從 tensordict 產生 state_dict。

除非將 flatten 設定為 True，否則 state-dict 的結構仍將是巢狀的。

tensordict 的 state-dict 包含重建 tensordict 所需的所有 tensors 和 meta-data (目前不支援名稱)。

參數:

destination (dict, optional) – 如果提供，tensordict 的狀態將會更新到這個 dict 中，並回傳同一個物件。否則，將會建立並回傳一個 OrderedDict。預設值：None。
prefix (str, optional) – 加到 tensor 名稱的前綴，用於組成 state_dict 中的鍵。預設值：''。
keep_vars (bool, optional) – 預設情況下，state dict 中回傳的 torch.Tensor 項目會從 autograd 中分離。如果設定為 True，則不會執行分離。預設值：False。
flatten (bool, optional) – 是否要使用 "." 字元來展平結構。預設為 False。

範例

>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, [])
>>> sd = data.state_dict()
>>> print(sd)
OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3', OrderedDict([('3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])), ('__batch_size', torch.Size([])), ('__device', None)])
>>> sd = data.state_dict(flatten=True)
OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3.3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])

std(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor¶

傳回輸入 tensordict 中所有元素的標準差值。

參數:

dim (int, int 的 tuple, optional) – 如果 None，則傳回一個無維度的 tensordict，其中包含所有葉節點的總和值（如果可以計算）。如果是整數或整數的 tuple，則只有在該維度與 tensordict 形狀相容時，才會在指定的維度上呼叫 std 。
keepdim (bool) – 輸出 tensor 是否保留 dim。

關鍵字參數:

correction (int) – 樣本大小與樣本自由度之間的差異。預設為 Bessel 的校正，correction=1。
reduce (bool, optional) – 如果 True，則 reduciton 將發生在所有 TensorDict 值上，並傳回單個簡化的 tensor。預設為 False。

sub(other: tensordict.base.TensorDictBase | float, *, alpha: Optional[float] = None, default: Optional[Union[str, Tensor]] = None)¶

從 self 減去 other，並將其縮放 alpha 倍。

\[\text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i\]

支援廣播 (broadcasting)、類型提升 (type promotion)，以及整數、浮點數和複數輸入。

參數:

other (TensorDict, Tensor 或 Number) – 要從 self 中減去的張量或數字。

關鍵字參數:

alpha (Number) – other 的乘數。
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。如果未提供任何值，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則只會考慮相交的鍵集，而其他鍵將被忽略。在所有其他情況下，default 將用於運算兩側的所有遺失條目。

sub_(other: tensordict.base.TensorDictBase | float, alpha: Optional[float] = None)¶: sub() 的原地 (in-place) 版本。

注意

原地 sub 不支援 default 關鍵字引數。

sum(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor¶

傳回輸入 tensordict 中所有元素的總和值。

參數:

dim (int, int 的 tuple, optional) – 如果為 None，則傳回一個無維度的 tensordict，其中包含所有葉節點的總和值 (如果可以計算)。如果是整數或整數元組，則僅當此維度與 tensordict 形狀相容時，才對指定的維度調用 sum 。
keepdim (bool) – 輸出 tensor 是否保留 dim。

關鍵字參數:

dtype (torch.dtype, optional) – 傳回的 tensor 所需的資料類型。如果指定，則在執行操作之前，輸入 tensor 會被轉換為 dtype。這對於防止資料類型溢位很有用。預設值： None。
reduce (bool, optional) – 如果 True，則 reduciton 將發生在所有 TensorDict 值上，並傳回單個簡化的 tensor。預設為 False。

tan() → T¶: 計算 TensorDict 中每個元素的 tan() 值。

tan_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 tan() 值。

tanh() → T¶: 計算 TensorDict 中每個元素的 tanh() 值。

tanh_() → T¶: 原地 (in-place) 計算 TensorDict 中每個元素的 tanh() 值。

to(*args, **kwargs) → TensorDictBase¶

將 TensorDictBase 子類別對應到另一個裝置、dtype 或另一個 TensorDictBase 子類別 (如果允許的話)。

不允許將 tensors 轉換為新的 dtype，因為 tensordicts 不一定只包含單一的 tensor dtype。

參數:

device (torch.device, optional) – tensordict 期望的裝置。
dtype (torch.dtype, optional) – tensordict 期望的浮點數或複數 dtype。
tensor (torch.Tensor, optional) – Tensor，其 dtype 和裝置是此 TensorDict 中所有 tensors 期望的 dtype 和裝置。

關鍵字參數:

non_blocking (bool, optional) – 操作是否應為 blocking。
memory_format (torch.memory_format, optional) – 此 tensordict 中 4D 參數和緩衝區期望的記憶體格式。
batch_size (torch.Size, optional) – 輸出 tensordict 的結果批次大小 (batch-size)。
other (TensorDictBase, optional) –
TensorDict 實例，其 dtype 和裝置是此 TensorDict 中所有 tensors 期望的 dtype 和裝置。

注意

由於 TensorDictBase 實例沒有 dtype，因此 dtype 是從範例 leaves 中收集的。如果有多個 dtype，則不會進行 dtype 轉換。
non_blocking_pin (bool, optional) –
如果 True，則張量會在傳送到裝置之前固定 (pinned)。這將以非同步方式完成，但可以透過 num_threads 引數進行控制。

注意

呼叫 tensordict.pin_memory().to("cuda") 通常會比 tensordict.to("cuda", non_blocking_pin=True) 慢得多，因為 pin_memory 在第二種情況下是以非同步方式呼叫的。如果張量很大且數量眾多，則多執行緒 pin_memory 通常會很有益：當要傳送的張量太少時，產生執行緒和收集資料的開銷會超過多執行緒的優勢，如果張量很小，則迭代長列表的開銷也會非常大。
num_threads (int or None, optional) – 如果 non_blocking_pin=True，則用於 pin_memory 的執行緒數。預設情況下，將產生 max(1, torch.get_num_threads()) 個執行緒。num_threads=0 將取消 pin_memory() 呼叫的任何多執行緒處理。

回傳值:

如果裝置與 tensordict 裝置不同，和/或如果傳遞了 dtype，則為新的 tensordict 實例。否則，為相同的 tensordict。batch_size 的修改僅在原地 (in-place) 完成。

注意

如果 TensorDict 已合併 (consolidated)，則產生的 TensorDict 也將被合併。每個新的張量都將是一個在合併儲存上投射到所需裝置的視圖。

範例

>>> data = TensorDict({"a": 1.0}, [], device=None)
>>> data_cuda = data.to("cuda:0")  # casts to cuda
>>> data_int = data.to(torch.int)  # casts to int
>>> data_cuda_int = data.to("cuda:0", torch.int)  # multiple casting
>>> data_cuda = data.to(torch.randn(3, device="cuda:0"))  # using an example tensor
>>> data_cuda = data.to(other=TensorDict({}, [], device="cuda:0"))  # using a tensordict example

to_dict(*, retain_none: bool = True) → dict[str, Any]¶

傳回一個字典，其鍵值對與 tensordict 的鍵值對匹配。

參數:: retain_none (bool) – 如果 True，則來自 tensorclass 實例的 None 值將被寫入字典中。否則，它們將被丟棄。預設值：True。

to_empty(*, device: Optional[Union[int, str, device]], recurse: bool = True) → T¶

將參數和緩衝區移動到指定的裝置，而不複製儲存空間。

參數:

device (torch.device) – 此模組中參數和緩衝區的所需裝置。
recurse (bool) – 子模組的參數和緩衝區是否應遞迴地移動到指定的裝置。

回傳值:

self

回傳類型:

模組

to_h5(filename, **kwargs)¶

將 tensordict 轉換為具有 h5 後端的 PersistentTensorDict。

參數:

filename (str or path) – h5 檔案的路徑。
device (torch.device or compatible, optional) – 張量返回後預期的裝置。預設為 None (預設在 cpu 上)。
**kwargs – 要傳遞給 h5py.File.create_dataset() 的 kwargs。

回傳值:

一個連結到新建立檔案的 PersitentTensorDict 實例。

範例

>>> import tempfile
>>> import timeit
>>>
>>> from tensordict import TensorDict, MemoryMappedTensor
>>> td = TensorDict({
...     "a": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000)),
...     "b": {"c": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000, 3))},
... }, [1_000_000])
>>>
>>> file = tempfile.NamedTemporaryFile()
>>> td_h5 = td.to_h5(file.name, compression="gzip", compression_opts=9)
>>> print(td_h5)
PersistentTensorDict(
    fields={
        a: Tensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float32, is_shared=False),
        b: PersistentTensorDict(
            fields={
                c: Tensor(shape=torch.Size([1000000, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([1000000]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([1000000]),
    device=None,
    is_shared=False)

to_module(module: Module, *, inplace: bool | None = None, return_swap: bool = True, swap_dest=None, use_state_dict: bool = False, non_blocking: bool = False, memo=None)¶

將 TensorDictBase 實例的內容遞迴地寫入給定的 nn.Module 屬性中。

to_module 也可以用作上下文管理器，以暫時使用參數/緩衝區集合填充模組（請參閱下面的示例）。

參數:

module (nn.Module) – 要將參數寫入的模組。

關鍵字參數:

inplace (bool, optional) – 如果 True，則模組中的參數或張量會就地更新。預設值為 False。
return_swap (bool, optional) – 如果 True，則會傳回舊的參數配置。預設值為 False。
swap_dest (TensorDictBase, optional) – 如果 return_swap 為 True，則為應該寫入交換內容的 tensordict。
use_state_dict (bool, optional) – 如果 True，將使用 state-dict API 來載入參數（包括 state-dict hooks）。預設值為 False。
non_blocking (bool, optional) – 如果 True 且此複製發生在不同裝置之間，則複製可能會相對於主機非同步發生。

範例

>>> from torch import nn
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = TensorDict.from_module(module)
>>> params.data.zero_()
>>> params.to_module(module)
>>> assert (module.layers[0].linear1.weight == 0).all()

將 tensordict 用作上下文管理器對於進行函數調用非常有用： .. rubric:: 範例

>>> from tensordict import from_module
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = TensorDict.from_module(module)
>>> params = params.data * 0 # Use TensorDictParams to remake these tensors regular nn.Parameter instances
>>> with params.to_module(module):
...     # Call the module with zeroed params
...     y = module(*inputs)
>>> # The module is repopulated with its original params
>>> assert (TensorDict.from_module(module) != 0).any()

回傳值:: 如果 return_swap 為 True，則為包含來自模組的值的 tensordict，否則為 None。

to_namedtuple(dest_cls: Optional[type] = None)¶

將 tensordict 轉換為 namedtuple。

參數:: dest_cls (Type, optional) – 一個可選的 namedtuple 類別來使用。

範例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> data.to_namedtuple()
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))

to_padded_tensor(padding=0.0, mask_key: Optional[NestedKey] = None)¶

將所有巢狀張量轉換為填充版本，並相應地調整批次大小。

參數:

padding (float) – tensordict 中張量的填充值。預設值為 0.0。
mask_key (NestedKey, optional) – 如果提供，則為要寫入有效值遮罩的鍵。如果異質維度不是 tensordict 批次大小的一部分，將導致錯誤。預設值為 None

to_pytree()¶

將 tensordict 轉換為 PyTree。

如果 tensordict 不是從 pytree 建立的，此方法只會傳回 self 而不進行修改。

請參閱 from_pytree() 以取得更多資訊和範例。

to_struct_array()¶

將 tensordict 轉換為 numpy 結構化陣列。

在一個 from_struct_array() - to_struct_array() 迴圈中，輸入和輸出陣列的內容應該一致。然而，to_struct_array 不會保留原始陣列的記憶體內容。

更多資訊請參考 from_struct_array()。

to_tensordict(*, retain_none: bool | None = None)¶

從 TensorDictBase 返回一個常規的 TensorDict 實例。

參數:

retain_none (bool) –

如果 True，tensorclass 實例中的 None 值將被寫入 tensordict 中。否則，它們將被丟棄。預設值：True。

注意

從 v0.8 開始，預設值將切換為 False。

回傳值:

一個包含相同數值的新 TensorDict 物件。

train(mode: bool = True) → T¶

將模組設定為訓練模式。

這僅對某些模組有效。請參閱特定模組的文件，了解它們在訓練/評估模式下的行為詳細資訊（如果它們受到影響），例如 Dropout、BatchNorm 等。

參數:: mode (bool) – 是否設定為訓練模式 (True) 或評估模式 (False)。預設值：True。
回傳值:: self
回傳類型:: 模組

transpose(dim0, dim1)¶

返回一個輸入的轉置版本的 tensordict。給定的維度 dim0 和 dim1 會被交換。

轉置 tensordict 的原地 (in-place) 或異地 (out-place) 修改也會影響原始 tensordict，因為記憶體是共享的，並且操作會映射回原始 tensordict。

範例

>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4])
>>> tensordict_transpose = tensordict.transpose(0, 1)
>>> print(tensordict_transpose.shape)
torch.Size([4, 3])
>>> tensordict_transpose.set("b",, torch.randn(4, 3))
>>> print(tensordict.get("b").shape)
torch.Size([3, 4])

trunc() → T¶: 計算 TensorDict 中每個元素的 trunc() 值。

trunc_() → T¶: 原地計算 TensorDict 中每個元素的 trunc() 值。

type(dst_type)¶

將所有 tensors 轉換為 dst_type。

參數:: dst_type (type 或 string) – 期望的類型

uint16()¶: 將所有 tensors 轉換為 torch.uint16。

uint32()¶: 將所有 tensors 轉換為 torch.uint32。

uint64()¶: 將所有 tensors 轉換為 torch.uint64。

uint8()¶: 將所有 tensors 轉換為 torch.uint8。

unbind(dim: int) → tuple[T, ...]¶

返回一個索引的 tensordicts 元組，沿指示的維度解綁。

範例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td0, td1, td2 = td.unbind(0)
>>> td0['x']
tensor([0, 1, 2, 3])
>>> td1['x']
tensor([4, 5, 6, 7])

unflatten(dim, unflattened_size)¶

解展平一個 tensordict 維度，將其擴展到期望的形狀。

參數:

dim (int) – 指定要解展平的輸入 tensor 的維度。
unflattened_size (shape) – 是 tensordict 解展平維度的新形狀。

範例

>>> td = TensorDict({
...     "a": torch.arange(60).view(3, 4, 5),
...     "b": torch.arange(12).view(3, 4)},
...     batch_size=[3, 4])
>>> td_flat = td.flatten(0, 1)
>>> td_unflat = td_flat.unflatten(0, [3, 4])
>>> assert (td == td_unflat).all()

unflatten_keys(separator: str = '.', inplace: bool = False) → TensorDictBase¶

將扁平化的 tensordict 遞迴地轉換為巢狀結構。

TensorDict 的類型將會遺失，結果會是一個簡單的 TensorDict 實例。巢狀 tensordict 的元數據將會從根目錄推斷：資料樹中所有實例將共享相同的批次大小、維度名稱和設備。

參數:

separator (str, optional) – 巢狀項目之間的分隔符。
inplace (bool, optional) – 如果 True，則產生的 tensordict 將與呼叫它的 tensordict 具有相同的 identity。預設為 False。

範例

>>> data = TensorDict({"a": 1, "b - c": 2, "e - f - g": 3}, batch_size=[])
>>> data.unflatten_keys(separator=" - ")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        e: TensorDict(
            fields={
                f: TensorDict(
                    fields={
                        g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
                    batch_size=torch.Size([]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

此方法和 unflatten_keys() 在處理 state-dicts 時特別有用，因為它們可以無縫地將扁平的字典轉換為模仿模型結構的資料結構。

範例

>>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4))
>>> ddp_model = torch.ao.quantization.QuantWrapper(model)
>>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".")
>>> print(state_dict)
TensorDict(
    fields={
        module: TensorDict(
            fields={
                0: TensorDict(
                    fields={
                        bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                        weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> model_state_dict = state_dict.get("module")
>>> print(model_state_dict)
TensorDict(
    fields={
        0: TensorDict(
            fields={
                bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> model.load_state_dict(dict(model_state_dict.flatten_keys(".")))

unlock_() → T¶

解除 tensordict 的鎖定，以便進行非原地 (non in-place) 操作。

可以作為裝飾器 (decorator) 使用。

詳情請參閱 lock_()。

unsqueeze(*args, **kwargs)¶

對介於 -td.batch_dims 和 td.batch_dims 之間的維度，將所有 tensors 進行 unsqueeze 操作，並將它們以新的 tensordict 形式返回。

參數:: dim (int) – 沿著此維度執行 unsqueeze 操作

範例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> td = td.unsqueeze(-2)
>>> td.shape
torch.Size([3, 1, 4])
>>> td.get("x").shape
torch.Size([3, 1, 4, 2])

此操作也可以作為上下文管理器 (context manager) 使用。對原始 tensordict 的變更將會以非原地 (out-place) 方式發生，也就是說，原始 tensors 的內容不會被修改。這也假設 tensordict 並未被鎖定（否則，需要先解鎖 tensordict）。

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> with td.unsqueeze(-2) as tds:
...     tds.set("y", torch.zeros(3, 1, 4))
>>> assert td.get("y").shape == [3, 4]

update(input_dict_or_td: dict[str, torch.Tensor] | tensordict.base.TensorDictBase, clone: bool = False, inplace: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None, is_leaf: Optional[Callable[[Type], bool]] = None) → TensorDictBase¶

使用字典或另一個 TensorDict 中的值來更新 TensorDict。

參數:

input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。
clone (bool, optional) – 是否在設定之前，先複製輸入的 (tensor) 字典中的 tensors。預設值為 False。
inplace (bool, optional) – 如果為 True，並且如果鍵與 tensordict 中現有的鍵匹配，則該鍵值對的更新將以原地 (in-place) 方式進行。如果找不到該條目，則將其新增。預設值為 False。

關鍵字參數:

keys_to_update (NestedKeys 的序列, optional) – 如果提供，則只會更新 key_to_update 中的鍵列表。這是為了避免呼叫 data_dest.update(data_src.select(*keys_to_update))。
non_blocking (bool, optional) – 如果 True 且此複製發生在不同裝置之間，則複製可能會相對於主機非同步發生。
is_leaf (Callable[[Type], bool], optional) – 一個可呼叫物件，用於判斷物件類型是否應被視為葉節點並進行交換，或者是一個張量集合（tensor collection）。可選參數。

回傳值:

self

範例

>>> td = TensorDict({}, batch_size=[3])
>>> a = torch.randn(3)
>>> b = torch.randn(3, 4)
>>> other_td = TensorDict({"a": a, "b": b}, batch_size=[])
>>> td.update(other_td, inplace=True) # writes "a" and "b" even though they can't be found
>>> assert td['a'] is other_td['a']
>>> other_td = other_td.clone().zero_()
>>> td.update(other_td)
>>> assert td['a'] is not other_td['a']

update_(input_dict_or_td: Union[dict[str, torch.Tensor], T], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) → T¶

使用字典或另一個 TensorDict 的值，就地更新 TensorDict。

與 update() 不同，如果 self 不知道鍵，此函數將拋出錯誤。

參數:

input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。
clone (bool, optional) – 是否在設定之前，先複製輸入的 (tensor) 字典中的 tensors。預設值為 False。

關鍵字參數:

keys_to_update (NestedKeys 序列, 可選) – 如果提供，則僅更新 key_to_update 中的鍵列表。這旨在避免調用 data_dest.update_(data_src.select(*keys_to_update))。
non_blocking (bool, optional) – 如果 True 且此複製發生在不同裝置之間，則複製可能會相對於主機非同步發生。

回傳值:

self

範例

>>> a = torch.randn(3)
>>> b = torch.randn(3, 4)
>>> td = TensorDict({"a": a, "b": b}, batch_size=[3])
>>> other_td = TensorDict({"a": a*0, "b": b*0}, batch_size=[])
>>> td.update_(other_td)
>>> assert td['a'] is not other_td['a']
>>> assert (td['a'] == other_td['a']).all()
>>> assert (td['a'] == 0).all()

update_at_(input_dict_or_td: Union[dict[str, torch.Tensor], T], idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) → T¶

使用來自字典或另一個 TensorDict 的值，在指定的索引處就地更新 TensorDict。

與 TensorDict.update 不同，如果 TensorDict 不知道鍵，此函數將拋出錯誤。

參數:

input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。
idx (int, torch.Tensor, iterable, slice) – 應該發生更新的 tensordict 的索引。
clone (bool, optional) – 輸入（tensor）字典中的 tensor 是否應該在設置之前被複製。預設值為 False。

關鍵字參數:

keys_to_update (NestedKeys 的序列, optional) – 如果提供，則只會更新 key_to_update 中的鍵列表。
non_blocking (bool, optional) – 如果 True 且此複製發生在不同裝置之間，則複製可能會相對於主機非同步發生。

回傳值:

self

範例

>>> td = TensorDict({
...     'a': torch.zeros(3, 4, 5),
...     'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4])
>>> td.update_at_(
...     TensorDict({
...         'a': torch.ones(1, 4, 5),
...         'b': torch.ones(1, 4, 10)}, batch_size=[1, 4]),
...    slice(1, 2))
TensorDict(
    fields={
        a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32),
        b: Tensor(torch.Size([3, 4, 10]), dtype=torch.float32)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> assert (td[1] == 1).all()

values(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) → Iterator[Tensor]¶

傳回代表 tensordict 值的產生器。

參數:

include_nested (bool, optional) – 若為 True，則會回傳巢狀 (nested) 的值。預設為 False。
leaves_only (bool, optional) – 若為 False，則只會回傳葉節點 (leaves)。預設為 False。
is_leaf – 一個可選的可呼叫 (callable) 物件，用於指示一個類別是否應被視為葉節點。

關鍵字參數:

sort (bool, optional) – 是否應該排序鍵。對於巢狀鍵，鍵會根據它們的連接名稱進行排序 (例如，("a", "key") 將被計為 "a.key" 以進行排序)。請注意，在處理大型 tensordict 時，排序可能會產生顯著的開銷。預設為 False。

var(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor¶

傳回輸入 tensordict 中所有元素的變異數值。

參數:

dim (int, int 的 tuple, optional) – 如果 None，則傳回包含所有葉節點總和值 (如果可以計算) 的無維度 tensordict。如果是整數或整數的 tuple，則僅當此維度與 tensordict 形狀相容時，才對指定的維度呼叫 var。
keepdim (bool) – 輸出 tensor 是否保留 dim。

關鍵字參數:

correction (int) – 樣本大小與樣本自由度之間的差異。預設為 Bessel 的校正，correction=1。
reduce (bool, optional) – 如果 True，則 reduciton 將發生在所有 TensorDict 值上，並傳回單個簡化的 tensor。預設為 False。

view(*shape: int, size: list | tuple | torch.Size | None = None, batch_size: torch.Size | None = None)¶

傳回一個 tensordict，其中包含根據與 tensordict batch_size 相容的新形狀所產生的張量視圖。

或者，可以提供 dtype 作為第一個未命名引數。在這種情況下，所有張量都將以相應的 dtype 進行檢視。請注意，這假設新形狀將與提供的 dtype 相容。有關 dtype 視圖的更多資訊，請參閱 view()。

參數:

*shape (int) – 結果 tensordict 的新形狀。
dtype (torch.dtype) – 或者，用於表示張量內容的 dtype。
size – iterable

關鍵字參數:

batch_size (torch.Size, optional) – 如果提供了 dtype，則可以使用此關鍵字引數重設 batch-size。如果使用形狀呼叫 view，則這無效。

回傳值:

具有所需 batch_size 的新 tensordict。

範例

>>> td = TensorDict(source={'a': torch.zeros(3,4,5),
...    'b': torch.zeros(3,4,10,1)}, batch_size=torch.Size([3, 4]))
>>> td_view = td.view(12)
>>> print(td_view.get("a").shape)  # torch.Size([12, 5])
>>> print(td_view.get("b").shape)  # torch.Size([12, 10, 1])
>>> td_view = td.view(-1, 4, 3)
>>> print(td_view.get("a").shape)  # torch.Size([1, 4, 3, 5])
>>> print(td_view.get("b").shape)  # torch.Size([1, 4, 3, 10, 1])

where(condition, other, *, out=None, pad=None)¶

傳回一個 TensorDict，其元素是根據條件從 self 或 other 中選取的。

參數:

condition (BoolTensor) – 當 True (非零) 時，產生 self，否則產生 other。
other (TensorDictBase 或 Scalar) – 在條件為 False 的索引處選取的值（如果 other 是純量）。

關鍵字參數:

out (TensorDictBase, optional) – 輸出的 TensorDictBase 實例。
pad (scalar, optional) – 如果提供此值，則來源或目標 tensordict 中缺少的鍵將會寫入為 torch.where(mask, self, pad) 或 torch.where(mask, pad, other)。預設值為 None，表示不容許遺失鍵。

xpu(device: Optional[Union[int, device]] = None) → T¶

將所有模型參數和緩衝區移動到 XPU。

這也會使相關的參數和緩衝區成為不同的物件。因此，如果模組將在 XPU 上運行並進行優化，則應在構建優化器之前調用它。

注意

此方法會就地修改模組。

參數:: device (int, optional) – 如果指定，所有參數將被複製到該裝置
回傳值:: self
回傳類型:: 模組

zero_() → T¶: 就地將 tensordict 中的所有張量歸零。

zero_grad(set_to_none: bool = True) → T¶

遞迴地將 TensorDict 的所有梯度歸零。

參數:: set_to_none (bool, optional) – 如果 True，tensor.grad 將為 None，否則為 0。預設值為 True。

TensorDictParams¶

文件

教學

資源