TensorDictParams¶

class tensordict.TensorDictParams(parameters: Optional[Union[TensorDictBase, dict]] = None, *, no_convert=False, lock: bool = False, **kwargs)¶

一個用於暴露引數的 TensorDict Base 類包裝器。

此類設計用於持有包含引數的 TensorDictBase 例項，使其可被父級 Module 訪問。這允許 tensordict 引數與 PyTorch 模組無縫整合，支援引數迭代和最佳化等操作。

主要特性

引數暴露：tensordict 內的引數會被暴露給父級模組，使其能夠包含在 named_parameters() 等操作中。
索引：索引的操作方式與被包裝的 tensordict 類似。但是，引數名稱（在 named_parameters() 中）會使用 TensorDict.flatten_keys(“_”) 進行註冊，這可能導致與 tensordict 內容不同的鍵名。
自動轉換：在 tensordict 中設定的任何張量都會自動轉換為 torch.nn.Parameter，除非透過 no_convert 關鍵字引數另行指定。

引數 (Args)

parameters (TensorDictBase or dict): 要表示為引數的 tensordict。值將被轉換為: parameters，除非設定了 no_convert=True。如果提供了 dict，它將被包裝在 TensorDict 例項中。也可以使用關鍵字引數。

關鍵字引數:

no_convert (bool) – 如果為 True，則不進行到 nn.Parameter 的轉換，所有非引數、非緩衝張量將被轉換為 Buffer 例項。如果為 False，則所有具有非整數 dtype 的張量將被轉換為 Parameter，而整數 dtype 將被轉換為 Buffer 例項。預設為 False。
lock (bool) –
如果為 True，則 TensorDictParams 託管的 tensordict 將被鎖定，阻止修改，並在需要 unlock_() 時可能影響效能。預設為 False。

警告

由於內部 tensordict 預設不會被複制或鎖定，因此在 TensorDictParams 中註冊 tensordict 並隨後修改其內容 __不會__ 更新 TensorDictParams parameters() 和 buffers() 序列中的值。
**kwargs – 用於填充 TensorDictParams 的鍵值對。與 parameters 輸入互斥。

示例

>>> 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 converted to Parameters
>>> m.params['other'] = torch.randn(3)
>>> assert isinstance(m.params['other'], nn.Parameter)

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

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

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

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

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

將 other 乘以 alpha 並加到 self 上。

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

引數:

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

關鍵字引數:

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

add_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor | float, *, alpha: float | None = None) → Any¶: 原地版本的 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) → Any¶

執行 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 (數字, 可選) – \(\text{tensor1} / \text{tensor2}\) 的乘數

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

addcmul(other1: tensordict.base.TensorDictBase | torch.Tensor, other2: tensordict.base.TensorDictBase | torch.Tensor, *, value: float | None = 1) → Any¶

執行 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 (數字, 可選) – \(other1 .* other2\) 的乘數

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

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

檢查 tensordict 中的所有值是否都為 True/非空。

引數:: dim (int, 可選) – 如果為 None，則返回一個布林值，指示所有張量是否都返回 tensor.all() == True。如果為整數，則當且僅當該維度與 tensordict 的形狀相容時，才會在指定的維度上呼叫 all。

amax(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None) → Self | torch.Tensor¶

返回輸入 tensordict 中所有元素的最大值。

與 max() 相同，但 return_indices=False。

amin(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None) → Self | torch.Tensor¶

返回輸入 tensordict 中所有元素的最小值。

與 min() 相同，但 return_indices=False。

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

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

引數:: dim (int, 可選) – 如果為 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 例項, 可選) – 如果提供，這些 tensordict 例項應具有與 self 匹配的結構。fn 引數應接收與 tensordicts 數量（包括 self）相同的未命名輸入。如果其他 tensordicts 存在缺失條目，可以透過 default 關鍵字引數傳遞預設值。

關鍵字引數:

batch_size (sequence of int, optional) – 如果提供，則結果 TensorDict 將具有所需的 batch_size。 batch_size 引數應與轉換後的 batch_size 匹配。這是一個僅限關鍵字的引數。
device (torch.device, 可選) – 生成的裝置，如果存在。
names (字串列表, 可選) – 新的維度名稱，以防batch_size被修改。
inplace (bool, optional) – 如果為 True，則原地修改。預設為 False。這是一個僅限關鍵字的引數。
default (Any, 可選) – 其他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 負責傳播其呼叫到巢狀級別。這允許在將呼叫傳播到巢狀 tensordicts 時進行精細控制。如果為 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, 可選) –

用於寫入結果的tensordict。這可以用來避免建立新的tensordict。

>>> td = TensorDict({"a": 0})
>>> td.apply(lambda x: x+1, out=td)
>>> assert (td==1).all()

警告

如果對 tensordict 執行的操作需要訪問多個鍵來進行單次計算，則將 out 引數設定為 self 可能會導致操作產生靜默錯誤的結果。例如

>>> 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建構函式的其他關鍵字引數。

返回:

一個包含轉換後張量的新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) → Any¶

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

引數:

fn (Callable) – 要應用於tensordict中張量的函式。
*others (TensorDictBase 序列, 可選) – 要使用的其他tensordicts。

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

返回:: self 或 self 的副本，其中應用了函式

as_tensor() → Any¶: 將 tensordict 的每個葉子節點轉換為普通的 torch.Tensor。

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

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

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

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

auto_batch_size_(batch_dims: Optional[int] = None, keep_compliant_size: bool = False) → Any¶

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

引數:

batch_dims (int, 可選) – 如果提供，則批次大小最多為 batch_dims。
keep_compliant_size (bool, 可選) – 如果為 True，則大小符合要求的子 tensordict 在 batch_dims 被傳遞時，其形狀不會改變。如果為 False，所有包含的 tensordicts 的 batch_dims 都將與 batch_dims 匹配。預設為 False。

返回:

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_() → Any¶

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

返回: self 及其編輯後的device屬性。

property batch_dims: int¶

tensordict批次大小的長度。

返回:: 描述tensordict維度的整數。

property batch_size: Size¶

TensorDict的形狀（或batch_size）。

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

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

`batch_size`可以動態編輯，如果新大小與TensorDict內容相容。例如，將批次大小設定為空值始終是允許的。

返回:: 一個 Size 物件，描述 TensorDict 的批次大小。

示例

>>> 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() → Any¶: 將所有張量轉換為torch.bfloat16。

bitwise_and(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

執行`self`和`other`之間的按位AND運算。

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

引數:: other (TensorDictBase 或 torch.Tensor) – 要執行按位 AND 的張量或 TensorDict。
關鍵字引數:: default (torch.Tensor 或 str, 可選) – 用於獨佔條目的預設值。如果未提供，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則僅考慮相交的鍵集，其他鍵將被忽略。在所有其他情況下，default 將用於操作兩側的所有缺失條目。

bool() → Any¶: 將所有張量轉換為torch.bool。

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

返回模組緩衝區的迭代器。

引數:: recurse (bool) – 如果為 True，則會生成此模組以及所有子模組的緩衝區。否則，僅生成此模組的直接成員緩衝區。
產生:: torch.Tensor – 模組緩衝區

示例

>>> # 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¶

計算包含的張量的位元組數。

關鍵字引數:: count_duplicates (bool) – 是否計算重複張量為獨立張量。如果為 False，則僅計算嚴格相同的張量（來自同一基張量的不同檢視但具有不同 id 的張量將被計算兩次）。預設為 True（每個張量被假定為單個副本）。

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

將tensordicts沿給定維度連線成一個tensordict。

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

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

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

引數:

dim (int, 可選) – 應沿哪個維度連線條目。

關鍵字引數:

sorted (bool 或 list of NestedKeys) – 如果為 True，則條目將按字母順序連線。如果為 False（預設），則使用字典順序。或者，可以提供一個鍵名列表，然後將張量按此順序連線。這會產生一些開銷，因為鍵列表將與 tensordict 中的葉節點名稱列表進行比較。
out (torch.Tensor, 可選) – 連線操作的可選目標張量。

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

將條目連線成一個新條目，並可能刪除原始值。

引數:

keys (NestedKey 序列) – 要連線的條目。

關鍵字引數:

out_key (NestedKey) – 連線輸入的新鍵名。
keep_entries (bool, optional) – 如果為 False，則 keys 中的條目將被刪除。預設為 False。
dim (int, optional) – 必須進行連線的維度。預設為 0。

返回: 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() → Any¶: 計算TensorDict中每個元素的ceil()值。

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

children() → Iterator[Module]¶

返回直接子模組的迭代器。

產生:: Module – 子模組

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

將tensordict分割成指定數量的塊，如果可能的話。

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

引數:

chunks (int) – 要返回的塊數
dim (int, 可選) – 分割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(min: tensordict.base.TensorDictBase | torch.Tensor = None, max: tensordict.base.TensorDictBase | torch.Tensor = None, *, out=None) → Any¶

將 self 中的所有元素限制在 [ min, max ] 範圍內。

令 min_value 和 max_value 分別為 min 和 max，則返回

\[y_i = \min(\max(x_i, \text{min\_value}_i), \text{max\_value}_i)\]

如果 min 為 None，則沒有下限。或者，如果 max 為 None，則沒有上限。

注意

如果 min 大於 max，則 torch.clamp(..., min, max) 將 input 中的所有元素設定為 max 的值。

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

如果`self`中的元素大於`other`，則將它們限制為`other`。

引數:: other (TensorDict 或 Tensor) – 另一個輸入tensordict或張量。
關鍵字引數:: default (torch.Tensor 或 str, 可選) – 用於獨佔條目的預設值。如果未提供，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則僅考慮相交的鍵集，其他鍵將被忽略。在所有其他情況下，default 將用於操作兩側的所有缺失條目。

clamp_max_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: 原地版本的 clamp_max()。

注意

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

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

將 self 的元素限制在不小於 other 的值。

引數:: other (TensorDict 或 Tensor) – 另一個輸入tensordict或張量。
關鍵字引數:: default (torch.Tensor 或 str, 可選) – 用於獨佔條目的預設值。如果未提供，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則僅考慮相交的鍵集，其他鍵將被忽略。在所有其他情況下，default 將用於操作兩側的所有缺失條目。

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

注意

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

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

clear_device_() → Any¶

清除 tensordict 的裝置。

返回: self

clear_refs_for_compile_() → Any¶

清除弱引用，以便 tensordict 安全地退出編譯區域。

在返回 TensorDict 之前遇到 torch._dynamo.exc.Unsupported: reconstruct: WeakRefVariable() 時，請使用此方法。

返回: self

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

將 TensorDictBase 子類例項克隆到相同型別的新 TensorDictBase 子類。

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

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

注意

與許多其他操作（逐元素算術、形狀操作等）不同，clone 不會繼承原始的鎖定屬性。此設計選擇是為了能夠建立一個可修改的副本，這是最常見的用法。

compile(*args, **kwargs)¶

使用 torch.compile() 編譯此 Module 的前向傳播。

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

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

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

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

complex64() → Any¶: 將所有張量轉換為 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) – 是否將合併後的資料放入 pinned 記憶體。預設為 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)¶: 返回一個相同型別的新 tensordict，其值是連續的（如果值已經是連續的，則返回 self）。(Returns a new tensordict of the same type with contiguous values (or self if values are already contiguous).)

copy() → Any¶

返回 tensordict 的淺複製（即，複製結構但不復制資料）。

等同於 TensorDictBase.clone(recurse=False)

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

請參閱 TensorDictBase.update_。

非阻塞引數將被忽略，僅為與 torch.Tensor.copy_() 的相容性而存在。

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

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

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

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

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

cpu()¶

將 tensordict 轉換為 CPU。

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

create_nested(key)¶

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

如果值已存在，它將被此操作覆蓋。此操作在鎖定的 tensordicts 中被阻止。

示例

>>> 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) – 如果提供，則為張量應被轉換到的 cuda 裝置。

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

cummax(dim: int, *, reduce: Optional[bool] = None, return_indices: bool = True) → Union[Any, Tensor]¶

返回輸入 tensordict 中所有元素的最大累積值。

引數:

dim (int) – 沿其執行 cummax 操作的維度的整數。

關鍵字引數:

reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。
return_argmins (bool, optional) – 當傳遞 dim 引數時，cummax() 返回一個包含值和索引的命名元組。此功能的 TensorDict 等價功能是返回一個具有相同結構的張量類，其中包含條目 "values" 和 "indices"。預設為 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummax(dim=0)
cummax(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummax(reduce=True, dim=0)
torch.return_types.cummax(...)

cummin(dim: int, *, reduce: Optional[bool] = None, return_indices: bool = True) → Union[Any, Tensor]¶

返回輸入 tensordict 中所有元素的最小累積值。

引數:

dim (int) – 沿其執行 cummin 操作的維度的整數。

關鍵字引數:

reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。
return_argmins (bool, optional) – 當傳遞 dim 引數時，cummin() 返回一個包含值和索引的命名元組。此功能的 TensorDict 等價功能是返回一個具有相同結構的張量類，其中包含條目 "values" 和 "indices"。預設為 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummin(dim=0)
cummin(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummin(reduce=True, dim=0)
torch.return_types.cummin(...)

property data¶: 返回一個包含葉張量的 .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 例項將顯示為巢狀的 tensordicts，以反映其葉節點的實際 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)¶

嘗試用連續張量（普通張量或巢狀張量）來表示懶惰堆疊。

關鍵字引數:: layout (torch.layout) – 巢狀張量的佈局，如果有的話。預設為 strided。

property depth: int¶

返回tensordict的深度 - 最大層級數。

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

detach() → Any¶

分離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, TensorCollection]] = None) → Any¶

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

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

支援廣播、型別提升以及整數、浮點數、tensordict或張量輸入。總是將整數型別提升為預設標量型別。

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

div_(other: tensordict.base.TensorDictBase | torch.Tensor) → Any¶: 原地執行的 div() 版本。

注意

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

double() → Any¶: 將所有張量轉換為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) → Any¶

將tensordict儲存到磁碟。

此函式是 memmap() 的代理。

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

返回一個新的、空的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)¶

返回條目的類，可能避免呼叫isinstance(td.get(key), type)。

當 get() 執行可能昂貴時，此方法應優於 tensordict.get(key).shape。

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

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

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

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

eval() → Self¶

將模組設定為評估模式。

這僅對某些模組有影響。有關模組在訓練/評估模式下的行為，例如它們是否受影響（如 Dropout、BatchNorm 等），請參閱具體模組的文件。

這等同於 self.train(False)。

有關 .eval() 和幾種可能與之混淆的類似機制之間的比較，請參閱區域性停用梯度計算。

返回:: self
返回型別:: 模組

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

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

值不會被複制：對原始tensordict或新tensordict的張量的就地修改將導致兩個tensordict都發生變化。

引數:

*keys (str) – 要排除的鍵。
inplace (bool) – 如果為 True，則 tensordict 將就地剪枝。預設為 False。

返回:

一個新的tensordict（如果inplace=True則為相同的tensordict），不包含被排除的條目。

示例

>>> 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() → Any¶: 計算TensorDict中每個元素的exp()值。

exp_() → Any¶: 就地計算TensorDict中每個元素的exp()值。

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

根據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._tensorcollection.TensorCollection | torch.Tensor) → Any¶

將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() → Any¶: 計算TensorDict中每個元素的expm1()值。

expm1_() → Any¶: 就地計算TensorDict中每個元素的expm1()值。

extra_repr() → str¶

返回模組的額外表示。

要列印自定義額外資訊，您應該在自己的模組中重新實現此方法。單行和多行字串均可接受。

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

Fills a tensor pointed by the key with a given scalar value.

引數:

key (str or nested key) – entry to be filled.
value (Number or bool) – value to use for the filling.

返回:

self

filter_empty_()¶: 就地過濾掉所有空的tensordict。

filter_non_tensor_data() → Any¶: 過濾掉所有非張量資料。

flatten(start_dim: int | None = None, end_dim: int | None = None)¶

展平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 具有相同的身份。預設為 False。
is_leaf (callable, optional) –
一個作用於類型別的可呼叫物件，返回一個布林值，指示該類是否應被視為葉子節點。

注意

is_leaf 的目的是不是阻止遞迴呼叫巢狀的 tensordicts，而是為過濾標記某些型別，當 `leaves_only=True` 時。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，巢狀 tensordict 的結構仍將被遍歷。換句話說，`is_leaf` 不控制遞迴深度，而是提供了一種當 `leaves_only=True` 時從結果中過濾掉某些型別的方法。這意味著樹中的一個節點既可以是葉節點，也可以是具有子節點的節點。實際上，`is_leaf` 的預設值不包含 tensordict 和 tensorclass 例項作為葉節點。

另請參閱

is_leaf_nontensor() 和 default_is_leaf()。

示例

>>> 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() → Any¶: 將所有張量轉換為torch.float。

float16() → Any¶: 將所有張量轉換為torch.float16。

float32() → Any¶: 將所有張量轉換為torch.float32。

float64() → Any¶: 將所有張量轉換為torch.float64。

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

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

forward(*input: Any) → None¶

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

所有子類都應重寫此方法。

注意

儘管前向傳播的實現需要在此函式中定義，但您應該在之後呼叫 Module 例項而不是此函式，因為前者會處理註冊的鉤子，而後者則會靜默忽略它們。

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

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

classmethod from_any(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

Recursively converts any object to a TensorDict.

注意

from_any 比常規的 TensorDict 建構函式限制更少。它可以使用自定義啟發式方法將資料結構（如 dataclasses 或 tuples）轉換為 tensordict。此方法可能會產生一些額外的開銷，並在對映策略方面涉及更多主觀選擇。

注意

This method recursively converts the input object to a TensorDict. If the object is already a TensorDict (or any similar tensor collection object), it will be returned as is.

引數:

obj – The object to be converted.

關鍵字引數:

auto_batch_size (bool, optional) – 如果為 True，將自動計算批次大小。預設為 False。
batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 的維度數。預設為 None（每個級別都有完整的批次大小）。
device (torch.device, optional) – 將建立 TensorDict 的裝置。
batch_size (torch.Size, optional) – TensorDict 的批次大小。與 auto_batch_size 互斥。

返回:

A TensorDict representation of the input object.

Supported objects

透過 from_dataclass() 的資料類（資料類將被轉換為 TensorDict 例項，而不是 tensorclasses）。
透過 from_namedtuple() 的命名元組。
透過 from_dict() 的字典。
透過 from_tuple() 的元組。
透過 from_struct_array() 的 NumPy 結構化陣列。
透過 from_h5() 的 HDF5 物件。

classmethod from_dataclass(dataclass, *, dest_cls: Optional[Type] = None, auto_batch_size: bool = False, batch_dims: Optional[int] = None, as_tensorclass: bool = False, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

Converts a dataclass into a TensorDict instance.

引數:

dataclass – The dataclass instance to be converted.

關鍵字引數:

dest_cls (tensorclass, optional) – 用於對映資料的 tensorclass 型別。如果未提供，則建立一個新類。如果 `obj` 是一個型別或 as_tensorclass 為 `False`，則無效。
auto_batch_size (bool, optional) – 如果為 True，則自動確定並應用批次大小到生成的 TensorDict。預設為 False。
batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 的維度數。預設為 None（每個級別都有完整的批次大小）。
as_tensorclass (bool, optional) – 如果為 True，則將轉換委託給函式 from_dataclass()，並返回一個相容張量的類（tensorclass()）或例項，而不是 TensorDict。預設為 False。
device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。

返回:

A TensorDict instance derived from the provided dataclass, unless as_tensorclass is True, in which case a tensor-compatible class or instance is returned.

丟擲:

TypeError – If the provided input is not a dataclass instance.

警告

This method is distinct from the free function from_dataclass and serves a different purpose. While the free function returns a tensor-compatible class or instance, this method returns a TensorDict instance.

注意

此方法建立一個新的 TensorDict 例項，其鍵對應於輸入 dataclass 的欄位。
結果 TensorDict 中的每個鍵都使用 `cls.from_any` 方法進行初始化。
auto_batch_size 選項允許自動確定批次大小並將其應用於結果 TensorDict。

classmethod from_dict(*args, **kwargs)¶

從字典或另一個 TensorDict 建立 TensorDict。

If batch_size is not specified, returns the maximum batch size possible.

This function works on nested dictionaries too, or can be used to determine the batch-size of a nested tensordict.

引數:

input_dict (dictionary, optional) – a dictionary to use as a data source (nested keys compatible).

關鍵字引數:

auto_batch_size (bool, optional) – 如果為 True，將自動計算批次大小。預設為 False。
batch_size (iterable of int, optional) – a batch size for the tensordict.
device (torch.device 或 相容型別, optional) – TensorDict 的裝置。
batch_dims (int, optional) – batch_dims（即用於 batch_size 的前導維度數量）。與 batch_size 互斥。注意，這是 tensordict 的__最大__批次維度數，較小的數字是允許的。
names (list of str, optional) – the dimension names of the 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, *, auto_batch_size: bool = False, 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: str = 'r', auto_batch_size: bool = False, batch_dims: Optional[int] = None, batch_size: Optional[Size] = None)¶

從 h5 檔案建立 PersistentTensorDict。

引數:

filename (str) – h5 檔案的路徑。

關鍵字引數:

mode (str, 可選) – 讀取模式。預設為 "r"。
auto_batch_size (bool, optional) – 如果為 True，將自動計算批次大小。預設為 False。
batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 的維度數。預設為 None（每個級別都有完整的批次大小）。
batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。

返回:

輸入 h5 檔案的 PersistentTensorDict 表示。

示例

>>> td = TensorDict.from_h5("path/to/file.h5")
>>> print(td)
PersistentTensorDict(
    fields={
        key1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        key2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

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 hook 時尤其有用。

示例

>>> 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 的 ensemable 學習/特徵期望應用檢索多個模組的引數。

引數:

modules (nn.Module 序列) – 要從中獲取引數的模組。如果模組結構不同，則需要懶惰堆疊（請參閱下面的 `lazy_stack` 引數）。

關鍵字引數:

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 hook 時尤其有用。
lazy_stack (bool, optional) –
是否密集堆疊或懶惰堆疊引數。預設為 False（密集堆疊）。

注意

lazy_stack 和 as_module 是互斥的特性。

警告

懶惰輸出和非懶惰輸出之間有一個重要的區別：非懶惰輸出將使用所需的批次大小重新例項化引數，而 `lazy_stack` 將僅將引數表示為懶惰堆疊。這意味著，雖然原始引數可以安全地傳遞給最佳化器（當 `lazy_stack=True` 時），但在設定為 `True` 時需要傳遞新引數。

警告

雖然使用 lazy stack 來保留原始引數引用很有誘惑力，但請記住，lazy stack 在每次呼叫 get() 時都會執行堆疊操作。這將需要記憶體（引數大小的 N 倍，如果構建了圖則更多）和計算時間。這也意味著最佳化器將包含更多引數，並且類似 step() 或 zero_grad() 的操作執行時間會更長。通常，lazy_stack 應僅用於少數用例。
expand_identical (bool, optional) – 如果為 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, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

遞迴地將命名元組轉換為 TensorDict。

引數:

named_tuple – 要轉換的命名元組例項。

關鍵字引數:

auto_batch_size (bool, optional) – 如果為 True，將自動計算批次大小。預設為 False。
batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 的維度數。預設為 None（每個級別都有完整的批次大小）。
device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。

返回:

輸入命名元組的 TensorDict 表示。

示例

>>> 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 到 tensordict 的雙射轉換 API。

當前接受的類包括列表、元組、命名元組和字典。

注意

對於字典，非 `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_remote_init(src: int, group: 'ProcessGroup' | None = None, device: torch.device | None = None) → Self¶

從遠端傳送的元資料建立新的 tensordict 例項。

此類方法接收由 init_remote 傳送的元資料，建立具有匹配形狀和 dtype 的新 tensordict，然後非同步接收實際的 tensordict 內容。

引數:

src (int) – 傳送元資料的源程序的秩。
group ("ProcessGroup", optional) – 要使用的程序組。預設為 None。
device (torch.device, 可選) – 用於張量運算的裝置。預設為 None。

返回:

使用接收到的元資料和內容初始化的新 tensordict 例項。

返回型別:

TensorDict

另請參閱

傳送程序應已呼叫 ~.init_remote 來發送元資料和內容。

classmethod from_struct_array(struct_array: ndarray, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None) → Any¶

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

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

注意

此方法執行零複製操作，這意味著生成的 TensorDict 將與輸入的 numpy 陣列共享相同的記憶體內容。因此，原地更改 numpy 陣列的值會影響 TensorDict 的內容。

引數:

struct_array (np.ndarray) – 要轉換的結構化 numpy 陣列。

關鍵字引數:

auto_batch_size (bool, optional) – 如果為 True，將自動計算批次大小。預設為 False。
batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 的維度數。預設為 None（每個級別都有完整的批次大小）。
device (torch.device, 可選) –
將建立 TensorDict 的裝置。預設為 None。

注意

更改裝置（即，指定任何非 `None` 或 `"cpu"` 的裝置）將傳輸資料，從而導致返回資料的記憶體位置發生更改。
batch_size (torch.Size, 可選) – TensorDict 的批次大小。預設為 None。

返回:

輸入的結構化 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 from_tuple(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

將元組轉換為 TensorDict。

引數:

obj – 要轉換的元組例項。

關鍵字引數:

auto_batch_size (bool, optional) – 如果為 True，將自動計算批次大小。預設為 False。
batch_dims (int, optional) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 的維度數。預設為 None（每個級別都有完整的批次大小）。
device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。

返回:

輸入的元組的 TensorDict 表示。

示例

>>> my_tuple = (1, 2, 3)
>>> td = TensorDict.from_tuple(my_tuple)
>>> print(td)
TensorDict(
    fields={
        0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

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

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

引數:

keys (list of NestedKey) – 指定新字典鍵的可迭代物件。
value (compatible type, optional) – 所有鍵的值。預設為 0。

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

沿由 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: 'torch.distributed.ProcessGroup' | None = None) → Self | None¶

從各個工作節點收集 tensordicts 並將它們堆疊到目標節點上的 self 中。

引數:

dst (int) – 目標工作節點的 rank，gather_and_stack() 將在該節點上呼叫。
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://: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://: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) → Union[Tensor, TensorCollection]¶

獲取輸入鍵對應的儲存值。

引數:

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

警告

以前，如果 tensordict 中不存在某個鍵且未提供預設值，則會引發 `KeyError`。從 v0.7 開始，此行為已更改，並返回 `None` 值（符合 dict.get 的行為）。要採用舊行為，請設定環境變數 `export TD_GET_DEFAULTS_TO_NONE=’0’` 或呼叫 :func`~tensordict.set_get_defaults_to_none(False)`。

注意

當處理不規則張量時，可以將關鍵字引數傳遞給 get()。有關完整概述，請參閱 get()。

示例

>>> td = TensorDict({"x": 1}, batch_size=[])
>>> td.get("x")
tensor(1)
>>> td.get("y")
None

get_at(key: NestedKey, *args, **kwargs) → Union[Tensor, TensorCollection]¶

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

引數:

key (str, tuple of str) – 要檢索的鍵。
index (int, slice, torch.Tensor, 可迭代物件) – 張量的索引。
default (torch.Tensor) – 如果 tensordict 中不存在該鍵，則返回的預設值。

返回:

索引的張量。

示例

>>> 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 – 要查詢的 buffer 的完全限定字串名稱。（要指定完全限定字串，請參閱 get_submodule。）
返回:: 由 target 引用的緩衝區
返回型別:: torch.Tensor
丟擲:: AttributeError – 如果目標字串引用了無效路徑或解析為非 buffer 的內容

get_extra_state() → Any¶

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

如果需要儲存額外的狀態，請實現此函式以及對應的 set_extra_state()。構建模組的 state_dict() 時會呼叫此函式。

注意，為了保證 state_dict 的序列化工作正常，額外狀態應該是可被 pickle 的。我們僅為 Tensors 的序列化提供向後相容性保證；其他物件的序列化形式若發生變化，可能導致向後相容性中斷。

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

get_item_shape(key: NestedKey)¶: 返回條目的形狀，可能避免呼叫 get()。

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

獲取非張量值（如果存在），或者在找不到非張量值時返回 default。

此方法對張量/TensorDict 值具有魯棒性，這意味著如果收集到的值是常規張量，它也會被返回（儘管此方法帶有一些開銷，不應超出其自然範圍使用）。

有關如何在 tensordict 中設定非張量值的更多資訊，請參閱 `tensordict.TensorDictBase.set_non_tensor`。

引數:

key (NestedKey) – 非張量資料的儲存位置。
default (Any, optional) – 找不到鍵時要返回的值。

返回： `tensordict.tensorclass.NonTensorData` 的內容，: 或對應於 `key` 的條目（如果它不是 `tensordict.tensorclass.NonTensorData`）或 `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 – 要查詢的 Parameter 的完全限定字串名稱。（要指定完全限定字串，請參閱 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)。因此，對於簡單地檢查某個子模組是否存在，應始終使用 get_submodule。

引數:: target – 要查詢的子模組的完全限定字串名稱。（要指定完全限定字串，請參閱上面的示例。）
返回:: 由 target 引用的子模組
返回型別:: torch.nn.Module
丟擲:: AttributeError – 如果沿目標字串解析路徑的任何一點解析為不存在的屬性名或不是 nn.Module 例項的物件。

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

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

init_remote(dst: int, group: 'ProcessGroup' | None = None, device: torch.device | None = None) → None¶

透過傳送元資料和內容來初始化遠端 tensordict。

此方法將當前 tensordict 的元資料（形狀、dtype 等）傳送到指定的目標 rank（dst）。

然後非同步傳送實際的 tensordict 內容。

引數:

dst (int) – 目標程序的 rank。
group ("ProcessGroup", optional) – 要使用的程序組。預設為 None。
device (torch.device, 可選) – 用於張量運算的裝置。預設為 None。

另請參閱

接收程序應呼叫 ~.from_remote_init 或等效方法來接收並基於傳送的元資料初始化新的 tensordict。

示例

>>> import os
>>> import torch
>>> import torch.distributed as dist
>>> from tensordict import TensorDict, MemoryMappedTensor
>>> import multiprocessing as mp
>>>
>>> def server(queue):
...     # Set environment variables for distributed communication
...     os.environ["MASTER_ADDR"] = "localhost"
...     os.environ["MASTER_PORT"] = "29505"
...
...     # Initialize the distributed backend
...     dist.init_process_group("gloo", rank=0, world_size=2)
...
...     # Create a sample tensordict
...     td = (
...         TensorDict(
...             {
...                 ("a", "b"): torch.ones(2),
...                 "c": torch.ones(2),
...                 ("d", "e", "f"): MemoryMappedTensor.from_tensor(torch.ones(2, 2)),
...             },
...             [2],
...         )
...         .expand(1, 2)
...         .contiguous()
...     )
...
...     # Send the tensordict metadata and content to the client
...     td.init_remote(dst=1)
...
>>> def client(queue):
...     # Set environment variables for distributed communication
...     os.environ["MASTER_ADDR"] = "localhost"
...     os.environ["MASTER_PORT"] = "29505"
...
...     # Initialize the distributed backend
...     dist.init_process_group("gloo", rank=1, world_size=2)
...
...     # Receive the tensordict metadata and content from the server
...     received_td = TensorDict.from_remote_init(src=0)
...
...     # Verify that the received tensordict matches the expected structure and values
...     assert set(received_td.keys()) == {"a", "c", "d"}
...     assert (received_td == 1).all()
...
...     # Signal that the test has completed successfully
...     queue.put("yuppie")
>>>
>>> if __name__ == "__main__":
...     queue = mp.Queue(1)
...
...     # Create and start the server and client processes
...     main_worker = mp.Process(target=server, args=(queue,))
...     secondary_worker = mp.Process(target=client, args=(queue,))
...
...     main_worker.start()
...     secondary_worker.start()
...
...     try:
...         out = queue.get(timeout=10)  # Wait for the signal with a timeout
...         print(out)  # Should print "yuppie"
...     finally:
...         queue.close()
...         main_worker.join(timeout=10)
...         secondary_worker.join(timeout=10)

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

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

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

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

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

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

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

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

注意

此方法就地修改模組。

引數:: device (int, optional) – 如果指定，所有引數都將複製到該裝置
返回:: self
返回型別:: 模組

irecv(src: int, *, group: 'torch.distributed.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。注意，這將根據 TensorDict 中包含的張量數量進行遞增。
pseudo_rand (bool) – 如果為 True，則標籤序列將是偽隨機的，允許從不同節點發送多個數據而不會重疊。請注意，這些偽隨機數的生成成本很高（每秒 1e-5），這意味著它可能會減慢演算法的執行速度。此值必須與傳遞給 isend() 的值匹配。預設為 False。

返回:

如果 return_premature=True，則返回一個未來列表以等待直到 tensordict 更新。: upon until the tensordict is updated.

is_consolidated()¶: 檢查 TensorDict 是否具有合併的儲存。

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

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

is_floating_point() → bool¶: 檢查 tensordict 中的所有張量是否為浮點數。

property is_memmap: bool¶

檢查 tensordict 是否為記憶體對映。

如果 TensorDict 例項是記憶體對映的，它將被鎖定（條目不能重新命名、刪除或新增）。如果 TensorDict 是使用所有記憶體對映的張量建立的，這並不意味著 `is_memmap` 將返回 `True`（因為新張量可能記憶體對映，也可能不記憶體對映）。只有在呼叫 `tensordict.memmap_()` 時，tensordict 才會被視為記憶體對映。

對於 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: 'torch.distributed.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False, return_early: bool = False) → int | List['Work']¶

非同步傳送 tensordict 的內容。

引數:

dst (int) – 應將內容傳送到的目標工作程序的 rank。

關鍵字引數:

group (torch.distributed.ProcessGroup, 可選) – 如果設定，將使用指定的程序組進行通訊。否則，將使用預設程序組。預設為 `None`。
init_tag (int) – 用於標記張量的初始 init_tag。注意，此值將根據 TensorDict 中包含的張量數量進行遞增。
pseudo_rand (bool) – 如果為 True，則標籤序列將是偽隨機的，允許從不同節點發送多個數據而不會重疊。請注意，這些偽隨機數的生成成本很高（每秒 1e-5），這意味著它可能會減慢演算法的執行速度。預設為 False。
return_early (bool, optional) – 如果為 True，則返回一個 futures 列表，而不是最後一個傳送的張量的 tag。預設為 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://: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://: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() → Any¶

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

實數值在非 NaN、負無窮或無窮大時是有限的。複數值在其實部和虛部都有限時是有限的。

isnan() → Any¶

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

當複數的實部和/或虛部為 NaN 時，複數值被視為 NaN。

isneginf() → Any¶: 測試輸入中的每個元素是否為負無窮。

isposinf() → Any¶: 測試輸入中的每個元素是否為負無窮。

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

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

返回 tensordict 的鍵值對生成器。

引數:

include_nested (bool, 可選) – 如果為 `True`，則返回巢狀值。預設為 `False`。
leaves_only (bool, 可選) – 如果為 `False`，則僅返回葉子節點。預設為 `False`。
is_leaf (callable, optional) –
一個作用於類型別的可呼叫物件，返回一個布林值，指示該類是否應被視為葉子節點。

注意

is_leaf 的目的是不是阻止遞迴呼叫巢狀的 tensordicts，而是為過濾標記某些型別，當 `leaves_only=True` 時。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，巢狀 tensordict 的結構仍將被遍歷。換句話說，`is_leaf` 不控制遞迴深度，而是提供了一種當 `leaves_only=True` 時從結果中過濾掉某些型別的方法。這意味著樹中的一個節點既可以是葉節點，也可以是具有子節點的節點。實際上，`is_leaf` 的預設值不包含 tensordict 和 tensorclass 例項作為葉節點。

另請參閱

is_leaf_nontensor() 和 default_is_leaf()。

關鍵字引數:

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

keys(*args, **kwargs)¶

返回 tensordict 鍵的生成器。

警告

TensorDict 的 `keys()` 方法返回鍵的懶惰檢視。如果查詢了 `keys` 但未迭代，然後修改了 tensordict，則稍後迭代 `keys` 將返回新配置的鍵。

引數:

include_nested (bool, 可選) – 如果為 `True`，則返回巢狀值。預設為 `False`。
leaves_only (bool, 可選) – 如果為 `False`，則僅返回葉子節點。預設為 `False`。
is_leaf (callable, optional) –
一個作用於類型別的可呼叫物件，返回一個布林值，指示該類是否應被視為葉子節點。

注意

is_leaf 的目的是不是阻止遞迴呼叫巢狀的 tensordicts，而是為過濾標記某些型別，當 `leaves_only=True` 時。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，巢狀 tensordict 的結構仍將被遍歷。換句話說，`is_leaf` 不控制遞迴深度，而是提供了一種當 `leaves_only=True` 時從結果中過濾掉某些型別的方法。這意味著樹中的一個節點既可以是葉節點，也可以是具有子節點的節點。實際上，`is_leaf` 的預設值不包含 tensordict 和 tensorclass 例項作為葉節點。

另請參閱

is_leaf_nontensor() 和 default_is_leaf()。

關鍵字引數:

sort (bool, optional) – 是否應該對鍵進行排序。對於巢狀鍵，鍵將根據其連線的名稱進行排序（即，("a", "key") 將被計為 "a.key" 以進行排序）。請注意，當處理大型 tensordicts 時，排序可能會產生顯著的開銷。預設為 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: int = 0, *, out=None, **kwargs)¶

建立 TensorDicts 的懶惰堆疊。

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

lerp(end: tensordict._tensorcollection.TensorCollection | torch.Tensor, weight: tensordict._tensorcollection.TensorCollection | torch.Tensor | float) → Any¶

根據標量或張量 `weight` 對兩個張量 `start`（由 `self` 提供）和 `end` 進行線性插值。

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

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

引數:

end (TensorDict) – 包含結束點的 tensordict。
weight (TensorDict, 張量或float) – 插值公式的權重。

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

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

lgamma_() → Any¶: 原地計算 TensorDict 中每個元素的 lgamma() 值。

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

從磁碟載入 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) → Any¶

從磁碟載入記憶體對映的 tensordict。

引數:

prefix (str 或 資料夾路徑) – 應從中獲取已儲存 tensordict 的資料夾路徑。
device (torch.device 或 等效項, 可選) – 如果提供，資料將非同步轉換為該裝置。支援 `"meta"` 裝置，在這種情況下，資料不會被載入，而是建立一組空的 "meta" 張量。這對於在不實際開啟任何檔案的情況下了解模型大小和結構很有用。
non_blocking (bool, 可選) – 如果為 `True`，則在將張量載入到裝置後不會呼叫同步。預設為 `False`。
out (TensorDictBase, 可選) – 應將資料寫入其中的可選 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”裝置上載入，或者作為假張量載入。

示例

>>> 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)¶

在呼叫 load_memmap_ 的 tensordict 中載入記憶體對映 tensordict 的內容。

有關更多資訊，請參閱 `tensordict.TensorDictBase.load_memmap`。

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

將格式與此 tensordict 的 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 字典中的項分配給 tensordict 中相應的鍵，而不是就地複製到 tensordict 的當前張量中。當 False 時，將保留當前模組張量的屬性；當 True 時，將保留 state 字典中張量的屬性。預設為 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_() → Any¶

鎖定 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() → Any¶: 計算 TensorDict 中每個元素的 log() 值。

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

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

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

log1p_() → Any¶: 原地計算 TensorDict 每個元素的 log1p() 值。

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

log2_() → Any¶: 原地計算 TensorDict 每個元素的 log2() 值。

log_() → Any¶: 原地計算 TensorDict 每個元素的 log() 值。

logical_and(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

對 self 和 other 執行邏輯 AND 操作。

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

引數:: other (TensorDictBase or torch.Tensor) – 用於執行邏輯 AND 操作的張量或 TensorDict。
關鍵字引數:: default (torch.Tensor 或 str, 可選) – 用於獨佔條目的預設值。如果未提供，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則僅考慮相交的鍵集，其他鍵將被忽略。在所有其他情況下，default 將用於操作兩側的所有缺失條目。

logsumexp(dim=None, keepdim=False, *, out=None)¶

返回給定維度 dim 上輸入 tensordict 各行的指數對數之和。計算是數值穩定的。

如果 keepdim 為 True，則輸出張量的大小與輸入張量相同，但在 dim 維度上大小為 1。否則，dim 將被壓縮（參見 squeeze()），導致輸出張量比輸入張量少 1 個（或 len(dim) 個）維度。

引數:

dim (int or tuple of ints, optional) – 要約簡的維度或維度。如果為 None，則約簡 tensordict 的所有批次維度。
keepdim (bool) – 輸出 tensordict 是否保留維度。

關鍵字引數:

out (TensorDictBase, 可選) – 輸出 tensordict。

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

根據形狀和可選的 dtype 建立一個空的記憶體對映張量。

警告

此方法在設計上不是執行緒安全的。存在於多個節點上的記憶體對映 TensorDict 例項需要使用 memmap_refresh_() 方法進行更新。

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

引數:

key (NestedKey) – 要寫入的新條目的鍵。如果鍵已存在於 tensordict 中，將引發異常。
shape (torch.Size or equivalent, torch.Tensor for nested tensors) – 要寫入的張量的形狀。

關鍵字引數:

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

返回:

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

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

根據儲存、形狀和可選的 dtype 建立一個空的記憶體對映張量。

警告

此方法在設計上不是執行緒安全的。存在於多個節點上的記憶體對映 TensorDict 例項需要使用 memmap_refresh_() 方法進行更新。

注意

如果儲存具有關聯的檔名，則必須與新檔案的檔名匹配。如果儲存沒有檔名，但 tensordict 具有關聯的路徑，這將導致異常。

引數:

key (NestedKey) – 要寫入的新條目的鍵。如果鍵已存在於 tensordict 中，將引發異常。
storage (torch.UntypedStorage) – 用於新 MemoryMappedTensor 的儲存。必須是物理記憶體儲存。
shape (torch.Size or equivalent, torch.Tensor for nested tensors) – 要寫入的張量的形狀。

關鍵字引數:

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

返回:

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

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

根據張量建立一個空的記憶體對映張量。

警告

此方法在設計上不是執行緒安全的。存在於多個節點上的記憶體對映 TensorDict 例項需要使用 memmap_refresh_() 方法進行更新。

如果 copy_data 為 True（即儲存是共享的），此方法將始終複製儲存內容。

引數:

key (NestedKey) – 要寫入的新條目的鍵。如果鍵已存在於 tensordict 中，將引發異常。
tensor (torch.Tensor) – 要在物理記憶體中複製的張量。

關鍵字引數:

copy_data (bool, 可選) – 如果為 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），其中塊將是原始資料的零複製切片，可以幾乎零成本地傳遞給程序。這使得處理非常大的資料集（例如 TB 級別）的成本非常低。

引數:

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

關鍵字引數:

out (TensorDictBase, 可選) – 輸出的可選容器。其在提供的 dim 上的批次大小必須與 self.ndim 匹配。如果它被共享或記憶體對映（is_shared() 或 is_memmap() 返回 True），它將在遠端程序中被填充，避免資料向內傳輸。否則，來自 self 切片的將被髮送到程序，在當前程序中收集，然後就地寫入 out。
chunksize (int, 可選) – 每個資料塊的大小。chunksize 為 0 將沿所需維度解綁 tensordict，並在函式應用後重新堆疊它，而 chunksize>0 將分割 tensordict 並對結果列表呼叫 torch.cat()。如果未提供，則塊的數量等於工作程序的數量。對於非常大的 tensordict，這樣大的塊可能無法放入記憶體中進行操作，並且可能需要更多的塊來使操作在實踐中可行。此引數與 num_chunks 互斥。
num_chunks (int, 可選) – 將 tensordict 分割成的塊數。如果未提供，則塊的數量等於工作程序的數量。對於非常大的 tensordict，這樣大的塊可能無法放入記憶體中進行操作，並且可能需要更多的塊來使操作在實踐中可行。此引數與 chunksize 互斥。
pool (mp.Pool, optional) – 要用於執行作業的多程序池例項。如果未提供，將在 map 方法中建立一個池。
generator (torch.Generator, optional) –
用於播種的生成器。將從其生成一個基礎種子，並且池中的每個工作程序將使用提供的種子（加上一個唯一的整數 0 到 num_workers）進行播種。如果沒有提供生成器，將使用一個隨機整數作為種子。要與未播種的工作程序一起使用，應單獨建立一個池並將其直接傳遞給 map()。

注意

當提供低數值種子時應謹慎，因為這可能導致實驗之間的自相關。例如：如果請求 8 個工作執行緒且種子為 4，則工作執行緒的種子範圍將是 4 到 11。如果種子是 5，則工作執行緒的種子範圍將是 5 到 12。這兩個實驗將有 7 個種子的重疊，這可能對結果產生意外影響。

注意

為工作執行緒設定種子的目的是讓每個工作執行緒擁有不同的種子，而不是為了讓 map 方法的呼叫結果可重現。換句話說，兩個實驗可能並且很可能返回不同的結果，因為無法知道哪個工作執行緒會執行哪個作業。但是，我們可以確保每個工作執行緒都有不同的種子，並且每個工作執行緒上的偽隨機操作將是不相關的。
max_tasks_per_child (int, 可選) – 每個子程序拾取的最大任務數。預設為 None，即對任務數沒有限制。
worker_threads (int, 可選) – 工作程序的執行緒數。預設為 1。
index_with_generator (bool, 可選) – 如果為 True，則 tensordict 的分割/分塊將在查詢期間進行，從而節省初始化時間。請注意，chunk() 和 split() 比索引（在生成器中使用）效率更高，因此初始化時間的處理時間增益可能會對總執行時間產生負面影響。預設為 False。
pbar (bool, 可選) – 如果為 True，將顯示進度條。需要可用 tqdm。預設為 False。
mp_start_method (str, optional) – 多程序的啟動方法。如果未提供，將使用預設啟動方法。可接受的字串是 "fork" 和 "spawn"。請注意，使用 "fork" 啟動方法時，"cuda" 張量無法在程序之間共享。如果 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[[TensorCollection], TensorCollection | 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) → Iterator[T]¶

沿一個維度迭代地將函式對映到tensordict的切片。

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

此方法將透過將tensordict分塊為大小相等的tensordicts並將操作分派到所需數量的工作程序來應用於tensordict例項。它將逐個生成結果。

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

注意

此方法特別適用於處理儲存在磁碟上的大型資料集（例如記憶體對映的 tensordict），其中塊將是原始資料的零複製切片，可以幾乎零成本地傳遞給程序。這使得處理非常大的資料集（例如 TB 級別）的成本非常低。

注意

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

引數:

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

關鍵字引數:

shuffle (bool, 可選) – 索引是否應全域性混洗。如果為 True，每個批次將包含不連續的樣本。如果 index_with_generator=False 且 shuffle=True，將引發錯誤。預設為 False。
chunksize (int, 可選) – 每個資料塊的大小。chunksize 為 0 將沿所需維度解綁 tensordict，並在函式應用後重新堆疊它，而 chunksize>0 將分割 tensordict 並對結果列表呼叫 torch.cat()。如果未提供，則塊的數量等於工作程序的數量。對於非常大的 tensordict，這樣大的塊可能無法放入記憶體中進行操作，並且可能需要更多的塊來使操作在實踐中可行。此引數與 num_chunks 互斥。
num_chunks (int, 可選) – 將 tensordict 分割成的塊數。如果未提供，則塊的數量等於工作程序的數量。對於非常大的 tensordict，這樣大的塊可能無法放入記憶體中進行操作，並且可能需要更多的塊來使操作在實踐中可行。此引數與 chunksize 互斥。
pool (mp.Pool, optional) – 要用於執行作業的多程序池例項。如果未提供，將在 map 方法中建立一個池。
generator (torch.Generator, optional) –
用於播種的生成器。將從其生成一個基礎種子，並且池中的每個工作程序將使用提供的種子（加上一個唯一的整數 0 到 num_workers）進行播種。如果沒有提供生成器，將使用一個隨機整數作為種子。要與未播種的工作程序一起使用，應單獨建立一個池並將其直接傳遞給 map()。

注意

當提供低數值種子時應謹慎，因為這可能導致實驗之間的自相關。例如：如果請求 8 個工作執行緒且種子為 4，則工作執行緒的種子範圍將是 4 到 11。如果種子是 5，則工作執行緒的種子範圍將是 5 到 12。這兩個實驗將有 7 個種子的重疊，這可能對結果產生意外影響。

注意

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

注意

index_with_generator 的預設值對於 map_iter 和 map 是不同的，前者假定將拆分後的 TensorDict 儲存在記憶體中成本過高。
pbar (bool, 可選) – 如果為 True，將顯示進度條。需要可用 tqdm。預設為 False。
mp_start_method (str, optional) – 多程序的啟動方法。如果未提供，將使用預設啟動方法。可接受的字串是 "fork" 和 "spawn"。請注意，使用 "fork" 啟動方法時，"cuda" 張量無法在程序之間共享。如果 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 的非原地版本。

引數:

mask (boolean torch.Tensor) – 要填充的值的掩碼。形狀必須與 tensordict 的批次大小匹配。
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 的批次大小匹配。
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) → Any¶

遮蔽 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.]])

max(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None, return_indices: bool = True) → Self | torch.Tensor¶

返回輸入 tensordict 中所有元素的最大值。

引數:

dim (int, optional) – 如果為 None，則返回一個無維度的 tensordict，其中包含所有葉子的最大值（如果可以計算）。如果為整數，則僅當該維度與 tensordict 的形狀相容時，才會在指定的維度上呼叫 max。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。
return_argmins (bool, 可選) – 當 dim 引數被傳遞時，max() 返回一個包含值和索引的命名元組。此方法的 TensorDict 等價物是返回一個具有條目 "values" 和 "indices" 的 tensorclass，其內部結構相同。預設為 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.max(dim=0)
max(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.max()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.max(reduce=True)
tensor(3.2942)

maximum(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

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

引數:: other (TensorDict 或 Tensor) – 另一個輸入tensordict或張量。
關鍵字引數:: default (torch.Tensor 或 str, 可選) – 用於獨佔條目的預設值。如果未提供，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則僅考慮相交的鍵集，其他鍵將被忽略。在所有其他情況下，default 將用於操作兩側的所有缺失條目。

maximum_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: maximum_ (maximum() 的就地版本。

注意

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

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

嘗試使 TensorDicts 密集堆疊，並在需要時回退到懶惰堆疊。

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

mean(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

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

引數:

dim (int, int 元組, str, 可選) – 如果為 None，則返回一個無維度的 tensordict，其中包含所有葉子的平均值（如果可計算）。如果為整數或整數元組，則當且僅當此維度與 tensordict 的形狀相容時，才會在指定的維度上呼叫 mean。目前只允許 “feature” 字串。使用 dim=”feature” 將在所有特徵維度上進行歸約。如果 reduce=True，將返回一個形狀與 TensorDict 的批次大小相同的張量。否則，將返回一個具有與 self 相同結構的 Tensordict，其中特徵維度已歸約。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定，則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設：None。
reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.mean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.mean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.mean(reduce=True)
tensor(-0.0547)
>>> td.mean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.mean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.mean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

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) → Any¶

將所有張量寫入記憶體對映的 Tensor 中，並放入新的 tensordict。

引數:

prefix (str) – 記憶體對映張量將儲存的目錄字首。目錄樹結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False（預設值），並且 tensordict 中某項已是儲存在磁碟上的張量且關聯了檔案，但未按 prefix 儲存到正確位置，則會引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

num_threads (int, 可選) – 用於寫入記憶體對映張量的執行緒數。預設為 0。
return_early (bool, 可選) – 如果為 True 且 num_threads>0，則方法將返回一個 tensordict 的 future。
share_non_tensor (bool, 可選) – 如果為 True，則非張量資料將在程序之間共享，並且在單個節點內的任何工作程序上進行寫入操作（例如就地更新或設定）將更新其他所有工作程序上的值。如果非張量葉子的數量很高（例如，共享大量非張量資料），這可能會導致 OOM 或類似錯誤。預設為 False。
existsok (bool, 可選) – 如果為 False，則如果同一路徑下已存在張量，將引發異常。預設為 True。

然後，Tensordict 被鎖定，這意味著任何非就地寫入操作（例如重新命名、設定或刪除條目）都將引發異常。一旦 tensordict 被解鎖，記憶體對映屬性將變為 False，因為不能保證跨程序身份。

返回:: 返回一個新的 tensordict，其中張量儲存在磁碟上（如果 return_early=False），否則返回一個 TensorDictFuture 例項。

注意

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

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

將所有張量原地寫入相應的記憶體對映張量。

引數:

prefix (str) – 記憶體對映張量將儲存的目錄字首。目錄樹結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False（預設值），並且 tensordict 中某項已是儲存在磁碟上的張量且關聯了檔案，但未按 prefix 儲存到正確位置，則會引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

num_threads (int, 可選) – 用於寫入記憶體對映張量的執行緒數。預設為 0。
return_early (bool, optional) – 如果為 True 且 num_threads>0，則該方法將返回一個 tensordict 的 future。可以使用 future.result() 查詢結果 tensordict。
share_non_tensor (bool, 可選) – 如果為 True，則非張量資料將在程序之間共享，並且在單個節點內的任何工作程序上進行寫入操作（例如就地更新或設定）將更新其他所有工作程序上的值。如果非張量葉子的數量很高（例如，共享大量非張量資料），這可能會導致 OOM 或類似錯誤。預設為 False。
existsok (bool, 可選) – 如果為 False，則如果同一路徑下已存在張量，將引發異常。預設為 True。

然後，Tensordict 被鎖定，這意味著任何非就地寫入操作（例如重新命名、設定或刪除條目）都將引發異常。一旦 tensordict 被解鎖，記憶體對映屬性將變為 False，因為不能保證跨程序身份。

返回:: 如果 return_early=False，則返回 self，否則返回 TensorDictFuture 例項。

注意

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

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

建立一個無內容的記憶體對映 tensordict，其形狀與原始 tensordict 相同。

引數:

prefix (str) – 記憶體對映張量將儲存的目錄字首。目錄樹結構將模仿 tensordict 的結構。
copy_existing (bool) – 如果為 False（預設值），並且 tensordict 中某項已是儲存在磁碟上的張量且關聯了檔案，但未按 prefix 儲存到正確位置，則會引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

num_threads (int, 可選) – 用於寫入記憶體對映張量的執行緒數。預設為 0。
return_early (bool, 可選) – 如果為 True 且 num_threads>0，則方法將返回一個 tensordict 的 future。
share_non_tensor (bool, 可選) – 如果為 True，則非張量資料將在程序之間共享，並且在單個節點內的任何工作程序上進行寫入操作（例如就地更新或設定）將更新其他所有工作程序上的值。如果非張量葉子的數量很高（例如，共享大量非張量資料），這可能會導致 OOM 或類似錯誤。預設為 False。
existsok (bool, 可選) – 如果為 False，則如果同一路徑下已存在張量，將引發異常。預設為 True。

然後，Tensordict 被鎖定，這意味著任何非就地寫入操作（例如重新命名、設定或刪除條目）都將引發異常。一旦 tensordict 被解鎖，記憶體對映屬性將變為 False，因為不能保證跨程序身份。

返回:: 如果 return_early=False，則建立一個新的 TensorDict 例項，其中資料儲存為記憶體對映張量；否則，建立一個 TensorDictFuture 例項。

注意

這是將一組大型緩衝區寫入磁碟的推薦方法，因為 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_()¶

如果記憶體對映的 tensordict 具有 saved_path，則重新整理其內容。

如果沒有任何路徑與之關聯，此方法將引發異常。

min(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None, return_indices: bool = True) → Self | torch.Tensor¶

返回輸入 tensordict 中所有元素的最小值。

引數:

dim (int, optional) – 如果 None，則返回一個無維度的 tensordict，其中包含所有葉子的最小值（如果可以計算）。如果為整數，則僅當此維度與 tensordict 的形狀相容時，才會對指定的維度呼叫 min。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。
return_argmins (bool, 可選) – 當 dim 引數被傳遞時，min() 返回一個包含值和索引的命名元組。此方法的 TensorDict 等價物是返回一個具有條目 "values" 和 "indices" 的 tensorclass，其內部結構相同。預設為 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.min(dim=0)
min(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.min()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.min(reduce=True)
tensor(-2.9953)

minimum(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

計算 self 和 other 的逐元素最小值。

引數:: other (TensorDict 或 Tensor) – 另一個輸入tensordict或張量。
關鍵字引數:: default (torch.Tensor 或 str, 可選) – 用於獨佔條目的預設值。如果未提供，則兩個 tensordict 的鍵列表必須完全匹配。如果傳遞了 default="intersection"，則僅考慮相交的鍵集，其他鍵將被忽略。在所有其他情況下，default 將用於操作兩側的所有缺失條目。

minimum_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: minimum_ (minimum() 的就地版本。

注意

原地 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[device, int]] = None) → Self¶

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

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

注意

此方法就地修改模組。

引數:: device (int, optional) – 如果指定，所有引數都將複製到該裝置
返回:: self
返回型別:: 模組

mul(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

將 other 乘以 self。

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

支援廣播、型別提升以及整數、浮點數和複數輸入。

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

mul_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: mul_ (mul() 的就地版本。

注意

原地 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 中的（名稱，張量）對的函式。對於每個葉子，只使用其葉子名稱（不使用完整的 NestedKey）。
*others (TensorDictBase 例項, 可選) – 如果提供，這些 tensordict 例項應具有與 self 匹配的結構。fn 引數應接收與 tensordicts 數量（包括 self）相同的未命名輸入。如果其他 tensordicts 存在缺失條目，可以透過 default 關鍵字引數傳遞預設值。
nested_keys (bool, optional) – 如果為 True，則將使用葉子的完整路徑。預設為 False，即僅將最後一個字串傳遞給函式。
batch_size (sequence of int, optional) – 如果提供，則結果 TensorDict 將具有所需的 batch_size。 batch_size 引數應與轉換後的 batch_size 匹配。這是一個僅限關鍵字的引數。
device (torch.device, 可選) – 生成的裝置，如果存在。
names (字串列表, 可選) – 新的維度名稱，以防batch_size被修改。
inplace (bool, optional) – 如果為 True，則就地進行更改。預設為 False。這是一個僅關鍵字引數。
default (Any, 可選) – 其他tensordict中缺失條目的預設值。如果未提供，缺失的條目將引發KeyError。
filter_empty (bool, optional) – 如果為 True，則將過濾掉空的 tensordicts。這還將降低計算成本，因為不會建立和銷燬空的資料結構。為了向後相容，預設為 False。
propagate_lock (bool, optional) – 如果為 True，則鎖定的 tensordict 將產生另一個鎖定的 tensordict。預設為 False。

call_on_nested (bool, optional) –

如果為 True，則該函式將應用於第一級張量和容器（TensorDict 或 tensorclass）。在這種情況下，func 負責傳播其呼叫到巢狀級別。這允許在將呼叫傳播到巢狀 tensordicts 時進行精細控制。如果為 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, 可選) –

用於寫入結果的tensordict。這可以用來避免建立新的tensordict。

>>> td = TensorDict({"a": 0})
>>> td.apply(lambda x: x+1, out=td)
>>> assert (td==1).all()

警告

如果對 tensordict 執行的操作需要訪問多個鍵來進行單次計算，則將 out 引數設定為 self 可能會導致操作產生靜默錯誤的結果。例如

>>> 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建構函式的其他關鍵字引數。

返回:

一個包含轉換後張量的新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, torch.Tensor]]¶

返回模組緩衝區上的迭代器，同時生成緩衝區的名稱和緩衝區本身。

引數:

prefix (str) – 要新增到所有緩衝區名稱前的字首。
recurse (bool) – 如果為 True，則會產生此模組及其所有子模組的緩衝區。否則，僅產生直接屬於此模組的緩衝區。
remove_duplicate (bool, 可選) – 是否移除結果中的重複緩衝區。預設為 True。

產生:

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

示例

>>> # 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) – 包含名稱和子模組的元組

示例

>>> # 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) – 名稱和模組的元組

注意

重複的模組只返回一次。在以下示例中，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, torch.nn.parameter.Parameter]]¶

返回模組引數的迭代器，同時生成引數的名稱和引數本身。

引數:

prefix (str) – 要新增到所有引數名稱前的字首。
recurse (bool) – 如果為 True，則會產生此模組及其所有子模組的引數。否則，僅產生直接屬於此模組的引數。
remove_duplicate (bool, 可選) – 是否移除結果中的重複引數。預設為 True。

產生:

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

示例

>>> # 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], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

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

引數:

dim (int, tuple of int, optional) – 如果 None，則返回一個無維度的 tensordict，其中包含所有葉子的平均值（如果可以計算）。如果為整數或整數元組，則僅當此維度與 tensordict 的形狀相容時，才會對指定的維度呼叫 mean。目前僅允許 “feature” 字串。使用 dim=”feature” 將實現跨所有特徵維度的約簡。如果 reduce=True，則將返回一個形狀與 TensorDict 的 batch_size 相同的張量。否則，將返回一個具有與 self 相同的結構但特徵維度已約簡的新 tensordict。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定，則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設：None。
reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.nanmean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.nanmean(reduce=True)
tensor(-0.0547)
>>> td.nanmean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.nanmean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

nansum(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

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

引數:

dim (int, tuple of int, optional) – 如果 None，則返回一個無維度的 tensordict，其中包含所有葉子的總和（如果可以計算）。如果為整數或整數元組，則僅當此維度與 tensordict 的形狀相容時，才會對指定的維度呼叫 sum。目前僅允許 “feature” 字串。使用 dim=”feature” 將實現跨所有特徵維度的約簡。如果 reduce=True，則將返回一個形狀與 TensorDict 的 batch_size 相同的張量。否則，將返回一個具有與 self 相同的結構但特徵維度已約簡的新 tensordict。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定，則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設：None。
reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nansum(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.nansum()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.nansum(reduce=True)
tensor(-0.)
>>> td.nansum(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nansum(reduce=True, dim="feature")
tensor([[15., 15., 15., 15.],
        [15., 15., 15., 15.],
        [15., 15., 15., 15.]])
>>> td.nansum(reduce=True, dim=0)
tensor([[9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.]])

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

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

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

neg_() → Any¶: 原地計算 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, empty_lazy: bool = False)¶

返回一個大小為 size 的 TensorDict，其中包含空的張量。

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

引數:

size (int...) – 定義輸出張量形狀的整數列表、元組或 torch.Size。

關鍵字引數:

dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 所需返回 tensordict 的裝置。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄返回張量的操作。預設為 False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需佈局。預設值：torch.strided。
pin_memory (bool, optional) – 如果設定為 True，則返回的張量將被分配到 pinned 記憶體中。僅適用於 CPU 張量。預設為 False。
empty_lazy (bool, optional) – 如果為 True，則 lazy stacks 將被清空其內容。當 lazy stack 的內容在填充新 tensordict 時可能會發生變化時，這會很有用。此引數將被傳播到子 tensordict。預設為 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, empty_lazy: bool = False)¶

返回一個大小為 size 的 TensorDict，其中填充值為 1。

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

引數:

size (sequence of int) – 定義輸出張量形狀的整數列表、元組或 torch.Size。
fill_value (scalar) – 用於填充輸出 Tensor 的數值。

關鍵字引數:

dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 所需返回 tensordict 的裝置。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄返回張量的操作。預設為 False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需佈局。預設值：torch.strided。
pin_memory (bool, optional) – 如果設定為 True，則返回的張量將被分配到 pinned 記憶體中。僅適用於 CPU 張量。預設為 False。
empty_lazy (bool, optional) – 如果為 True，則 lazy stacks 將被清空其內容。當 lazy stack 的內容在填充新 tensordict 時可能會發生變化時，這會很有用。此引數將被傳播到子 tensordict。預設為 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, empty_lazy: bool = False)¶

返回一個大小為 size 的 TensorDict，其中填充值為 1。

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

引數:

size (int...) – 定義輸出張量形狀的整數列表、元組或 torch.Size。

關鍵字引數:

dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 所需返回 tensordict 的裝置。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄返回張量的操作。預設為 False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需佈局。預設值：torch.strided。
pin_memory (bool, optional) – 如果設定為 True，則返回的張量將被分配到 pinned 記憶體中。僅適用於 CPU 張量。預設為 False。
empty_lazy (bool, optional) – 如果為 True，則 lazy stacks 將被清空其內容。當 lazy stack 的內容在填充新 tensordict 時可能會發生變化時，這會很有用。此引數將被傳播到子 tensordict。預設為 False。

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

使用 data 作為張量返回一個新的 TensorDict。

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

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

引數:

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 應該記錄返回張量的操作。預設為 False。
pin_memory (bool, optional) – 如果設定為 True，則返回的張量將被分配到 pinned 記憶體中。僅適用於 CPU 張量。預設為 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, empty_lazy: bool = False)¶

返回一個大小為 size 的 TensorDict，其中填充值為 0。

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

引數:

size (int...) – 定義輸出張量形狀的整數列表、元組或 torch.Size。

關鍵字引數:

dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值：如果 None，則 torch.dtype 將保持不變。
device (torch.device, optional) – 所需返回 tensordict 的裝置。預設值：如果 None，則 torch.device 將保持不變。
requires_grad (bool, optional) – 如果 autograd 應該記錄返回張量的操作。預設為 False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需佈局。預設值：torch.strided。
pin_memory (bool, optional) – 如果設定為 True，則返回的張量將被分配到 pinned 記憶體中。僅適用於 CPU 張量。預設為 False。
empty_lazy (bool, optional) – 如果為 True，則 lazy stacks 將被清空其內容。當 lazy stack 的內容在填充新 tensordict 時可能會發生變化時，這會很有用。此引數將被傳播到子 tensordict。預設為 False。

non_tensor_items(include_nested: bool = False)¶: 遞迴地返回所有非張量葉節點。

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

計算 tensordict 中每個張量的範數。

關鍵字引數:

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

numel() → int¶

批次中的總元素數量。

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

numpy() → numpy.ndarray | dict[str, Any]¶

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

非張量資料按原樣公開。

示例

>>> 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!'}}

另請參閱 to_struct_array() 轉換為結構化陣列。

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

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

關鍵字引數:: count_duplicates (bool) – 是否計算重複張量為獨立張量。如果為 False，則僅計算嚴格相同的張量（來自同一基張量的不同檢視但具有不同 id 的張量將被計算兩次）。預設為 True（每個張量被假定為單個副本）。

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)¶

根據 dims 返回 tensordict 的檢視，其中批次維度已重新排序。

引數:

*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)¶

在儲存的張量上呼叫 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) → Any¶

在儲存的張量上呼叫 pin_memory()，並返回就地修改的 TensorDict。

引數:: num_threads (int 或 str) – 如果提供，則用於在葉子上呼叫 pin_memory 的執行緒數。如果傳遞 "auto"，則自動確定執行緒數。

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

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

如果未找到該值且未提供預設值，則會引發 KeyError。

引數:

key (str or nested key) – 要查詢的條目。
default (Any, optional) – 找不到鍵時要返回的值。

示例

>>> 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, TensorCollection]] = None) → Any¶

將 self 的每個元素與 other 求冪，並返回結果張量。

other 可以是單個 float 數、Tensor 或 TensorDict。

當 other 是張量時，input 和 other 的形狀必須是可廣播的。

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

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

注意

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

prod(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回輸入 TensorDict 中所有元素的乘積。

引數:

dim (int, tuple of int, optional) – 如果 None，則返回一個無維度的 tensordict，其中包含所有葉子的積（如果可以計算）。如果為整數或整數元組，則僅當此維度與 tensordict 的形狀相容時，才會對指定的維度呼叫 prod。目前僅允許 “feature” 字串。使用 dim=”feature” 將實現跨所有特徵維度的約簡。如果 reduce=True，則將返回一個形狀與 TensorDict 的 batch_size 相同的張量。否則，將返回一個具有與 self 相同的結構但特徵維度已約簡的新 tensordict。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定，則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設：None。
reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.prod(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.prod()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.prod(reduce=True)
tensor(-0.)
>>> td.prod(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.prod(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.prod(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

qint32() → Any¶: 將所有張量強制轉換為 torch.qint32。

qint8() → Any¶: 將所有張量強制轉換為 torch.qint8。

quint4x2() → Any¶: 將所有張量強制轉換為 torch.quint4x2。

quint8() → Any¶: 將所有張量強制轉換為 torch.quint8。

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

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

record_stream(stream: Stream) → Any¶

標記 TensorDict 已被此流使用。

當 TensorDict 被解除分配時，請確保在解除分配時該流上的所有工作完成之前，不要將張量記憶體重用於其他張量。

更多資訊請參閱 record_stream()。

recv(src: int, *, group: 'torch.distributed.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。注意，這將根據 TensorDict 中包含的張量數量進行遞增。
pseudo_rand (bool) – 如果為 True，則標籤序列將是偽隨機的，允許從不同節點發送多個數據而不會重疊。請注意，這些偽隨機數的生成成本很高（每秒 1e-5），這意味著它可能會減慢演算法的執行速度。此值必須與傳遞給 send() 的值匹配。預設為 False。

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

在所有機器上歸約 TensorDict。

只有具有 rank dst 的程序才能接收最終結果。

refine_names(*names) → Any¶

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

精煉是重新命名的特殊情況，它“提升”未命名維度。None 維度可以精煉為任何名稱；已命名維度只能精煉為相同名稱。

由於命名張量可以與未命名張量共存，因此細化名稱提供了一種很好的方式來編寫既適用於命名張量又適用於未命名張量的命名張量感知程式碼。

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

Returns: 具有根據輸入命名的維度的相同的 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[torch.Tensor, ...], Tensor], Union[tuple[torch.Tensor, ...], Tensor]], Union[None, tuple[torch.Tensor, ...], Tensor]]) → RemovableHandle¶

在模組上註冊一個反向傳播鉤子。

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

返回:: 一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子
返回型別:: torch.utils.hooks.RemovableHandle

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

向模組新增一個緩衝區。

這通常用於註冊一個不應被視為模型引數的 buffer。例如，BatchNorm 的 running_mean 不是一個引數，但它是 module 狀態的一部分。Buffer 預設是持久化的，並且會與引數一起儲存。透過將 persistent 設定為 False，可以改變這種行為。持久化 buffer 和非持久化 buffer 之間的唯一區別是，後者不會成為此 module 的 state_dict 的一部分。

可以使用給定名稱作為屬性訪問緩衝區。

引數:

name (str) – 緩衝區的名稱。可以使用給定名稱從此模組訪問緩衝區
tensor (Tensor or None) – 要註冊的 buffer。如果為 None，則對 buffer 執行的操作（如 cuda）將被忽略。如果為 None，則 buffer **不**包含在 module 的 state_dict 中。
persistent (bool) – buffer 是否是此 module 的 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。

如果 with_kwargs 為 False 或未指定，則輸入僅包含傳遞給 module 的位置引數。關鍵字引數不會傳遞給 hook，只會傳遞給 forward。hook 可以修改輸出。它可以就地修改輸入，但這不會影響 forward，因為這是在 forward() 被呼叫之後才執行。hook 應該具有以下簽名

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

如果 with_kwargs 為 True，則前向鉤子將接收傳遞給 forward 函式的 kwargs，並需要返回可能已修改的輸出。鉤子應該具有以下簽名

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

引數:

hook (Callable) – 使用者定義的待註冊鉤子。
prepend (bool) – 如果為 True，則提供的 hook 將在當前 torch.nn.Module 的所有現有 forward hook 之前觸發。否則，提供的 hook 將在當前 torch.nn.Module 的所有現有 forward hook 之後觸發。請注意，使用 register_module_forward_hook() 註冊的全域性 forward hook 將在為此方法註冊的所有 hook 之前觸發。預設值：False
with_kwargs (bool) – 如果為 True，則 hook 將接收傳遞給 forward 函式的 kwargs。預設值：False
always_call (bool) – 如果為 True，則無論在呼叫 Module 時是否引發異常，都會執行 hook。預設值：False

返回:

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別:

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() 之前，都會呼叫 hook。

如果 with_kwargs 為 false 或未指定，則輸入僅包含傳遞給模組的位置引數。關鍵字引數不會傳遞給鉤子，而只會傳遞給 forward。鉤子可以修改輸入。使用者可以返回一個元組或單個修改後的值。我們將把值包裝成一個元組，如果返回的是單個值（除非該值本身就是元組）。鉤子應該具有以下簽名

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

如果 with_kwargs 為 true，則前向預鉤子將接收傳遞給 forward 函式的 kwargs。如果鉤子修改了輸入，則應該返回 args 和 kwargs。鉤子應該具有以下簽名

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

引數:

hook (Callable) – 使用者定義的待註冊鉤子。
prepend (bool) – 如果為 True，則提供的 hook 將在當前 torch.nn.Module 的所有現有 forward_pre hook 之前觸發。否則，提供的 hook 將在當前 torch.nn.Module 的所有現有 forward_pre hook 之後觸發。請注意，使用 register_module_forward_pre_hook() 註冊的全域性 forward_pre hook 將在為此方法註冊的所有 hook 之前觸發。預設值：False
with_kwargs (bool) – 如果為 True，則 hook 將接收傳遞給 forward 函式的 kwargs。預設值：False

返回:

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別:

torch.utils.hooks.RemovableHandle

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

在模組上註冊一個反向傳播鉤子。

每次計算相對於模組的梯度時，將呼叫此鉤子，其觸發規則如下：

通常，鉤子在計算相對於模組輸入的梯度時觸發。

如果模組輸入都不需要梯度，則在計算相對於模組輸出的梯度時觸發鉤子。

如果模組輸出都不需要梯度，則鉤子將不觸發。

鉤子應具有以下簽名

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

關於輸入和輸出的梯度（grad_input 和 grad_output）是包含輸入的梯度和輸出的梯度的元組。Hook 不應該修改其引數，但它可以選擇性地返回一個關於輸入的新的梯度，該梯度將被用於替換後續計算中的 grad_input。grad_input 只對應於作為位置引數給出的輸入，並且所有關鍵字引數都會被忽略。對於所有非 Tensor 引數，grad_input 和 grad_output 中的條目將為 None。

由於技術原因，當此鉤子應用於模組時，其前向函式將接收傳遞給模組的每個張量的檢視。類似地，呼叫者將接收模組前向函式返回的每個張量的檢視。

警告

使用反向傳播鉤子時不允許就地修改輸入或輸出，否則將引發錯誤。

引數:

hook (Callable) – 要註冊的使用者定義鉤子。
prepend (bool) – 如果為 True，則提供的 hook 將在當前 torch.nn.Module 的所有現有 backward hook 之前觸發。否則，提供的 hook 將在當前 torch.nn.Module 的所有現有 backward hook 之後觸發。請注意，使用 register_module_full_backward_hook() 註冊的全域性 backward hook 將在為此方法註冊的所有 hook 之前觸發。

返回:

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別:

torch.utils.hooks.RemovableHandle

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

在模組上註冊一個反向預鉤子。

每次計算模組的梯度時，將呼叫此鉤子。鉤子應具有以下簽名

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

grad_output 是一個元組。鉤子不應修改其引數，但可以選擇返回一個新的輸出梯度，該梯度將取代 grad_output 用於後續計算。對於所有非 Tensor 引數，grad_output 中的條目將為 None。

由於技術原因，當此鉤子應用於模組時，其前向函式將接收傳遞給模組的每個張量的檢視。類似地，呼叫者將接收模組前向函式返回的每個張量的檢視。

警告

使用反向傳播鉤子時不允許就地修改輸入，否則將引發錯誤。

引數:

hook (Callable) – 要註冊的使用者定義鉤子。
prepend (bool) – If true, the provided hook will be fired before all existing backward_pre hooks on this torch.nn.Module. Otherwise, the provided hook will be fired after all existing backward_pre hooks on this torch.nn.Module. Note that global backward_pre hooks registered with register_module_full_backward_pre_hook() will fire before all hooks registered by this method.

返回:

一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子

返回型別:

torch.utils.hooks.RemovableHandle

register_get_post_hook(hook)¶: Register a hook to be called after any get operation on leaf tensors.

register_load_state_dict_post_hook(hook)¶

註冊一個後鉤子，用於在模組的 load_state_dict() 被呼叫後執行。

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

The module argument is the current module that this hook is registered on, and the incompatible_keys argument is a NamedTuple consisting of attributes missing_keys and unexpected_keys. missing_keys is a list of str containing the missing keys and unexpected_keys is a list of str containing the unexpected keys.

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

Note that the checks performed when calling load_state_dict() with strict=True are affected by modifications the hook makes to missing_keys or unexpected_keys, as expected. Additions to either set of keys will result in an error being thrown when strict=True, and clearing out both missing and unexpected keys will avoid an error.

返回:: 一個控制代碼，可用於透過呼叫 handle.remove() 來移除新增的鉤子
返回型別:: torch.utils.hooks.RemovableHandle

register_load_state_dict_pre_hook(hook)¶

註冊一個預鉤子，用於在模組的 load_state_dict() 被呼叫之前執行。

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

引數:: hook (Callable) – 在載入狀態字典之前將呼叫的可呼叫鉤子。

register_module(name: str, module: Optional[Module]) → None¶: Alias for add_module().

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

向模組新增一個引數。

可以使用給定名稱作為屬性訪問該引數。

引數:

name (str) – 引數的名稱。可以使用給定名稱從此模組訪問該引數
param (Parameter or None) – parameter to be added to the module. If None, then operations that run on parameters, such as cuda, are ignored. If None, the parameter is not included in the module’s state_dict.

register_state_dict_post_hook(hook)¶

註冊 state_dict() 方法的後置鉤子。

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

註冊的鉤子可以就地修改 state_dict。

register_state_dict_pre_hook(hook)¶

註冊 state_dict() 方法的前置鉤子。

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

註冊的鉤子可用於在進行 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)¶

Same as rename(), but executes the renaming in-place.

示例

>>> 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

repeat(*repeats: int) → Any¶

沿指定維度重複此張量。

Unlike expand(), this function copies the tensor’s data.

警告

repeat() behaves differently from repeat(), but is more similar to numpy.tile(). For the operator similar to numpy.repeat(), see repeat_interleave().

引數:: repeat (torch.Size, int..., tuple of int or list of int) – 沿每個維度重複此張量的次數。

示例

>>> import torch
>>>
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     {
...         "a": torch.randn(3, 4, 5),
...         "b": TensorDict({
...             "c": torch.randn(3, 4, 10, 1),
...             "a string": "a string!",
...         }, batch_size=[3, 4, 10])
...     }, batch_size=[3, 4],
... )
>>> print(td.repeat(1, 2))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 8, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                a string: NonTensorData(data=a string!, batch_size=torch.Size([3, 8, 10]), device=None),
                c: Tensor(shape=torch.Size([3, 8, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 8, 10]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 8]),
    device=None,
    is_shared=False)

repeat_interleave(*shape: int)¶

重複 TensorDict 的元素。

警告

這與 repeat() 不同，但與 numpy.repeat() 類似。

引數:

repeats (torch.Tensor or int) – 每個元素的重複次數。repeats 將被廣播以適應給定軸的形狀。
dim (int, optional) – 重複值的維度。預設為使用展平的輸入陣列，並返回一個展平的輸出陣列。

關鍵字引數:

output_size (int, optional) – Total output size for the given axis (e.g. sum of repeats). If given, it will avoid stream synchronization needed to calculate output shape of the tensordict.

返回:

重複的 TensorDict，其形狀與輸入相同，除了在給定軸上。

示例

>>> import torch
>>>
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     {
...         "a": torch.randn(3, 4, 5),
...         "b": TensorDict({
...             "c": torch.randn(3, 4, 10, 1),
...             "a string": "a string!",
...         }, batch_size=[3, 4, 10])
...     }, batch_size=[3, 4],
... )
>>> print(td.repeat_interleave(2, dim=0))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([6, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                a string: NonTensorData(data=a string!, batch_size=torch.Size([6, 4, 10]), device=None),
                c: Tensor(shape=torch.Size([6, 4, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([6, 4, 10]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([6, 4]),
    device=None,
    is_shared=False)

replace(*args, **kwargs)¶

建立 TensorDict 的淺複製，其中條目已被替換。

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

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

requires_grad_(requires_grad=True) → Any¶

更改 autograd 是否應記錄此張量上的操作：就地設定此張量的 requires_grad 屬性。

返回此 TensorDict。

引數:: requires_grad (bool, optional) – autograd 是否應該記錄此 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() → Any¶: 計算 TensorDict 中每個元素的 round() 值。

round_() → Any¶: 原地計算 TensorDict 中每個元素的 round() 值。

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

從 self 中減去 other 乘以 alpha 的縮放值。

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

支援廣播、型別提升以及整數、浮點數和複數輸入。

引數:

other (TensorDict, Tensor 或 Number) – 從 self 中減去的張量或數字。

關鍵字引數:

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

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

將tensordict儲存到磁碟。

此函式是 memmap() 的代理。

property saved_path¶

返回 memmap 儲存的 TensorDict 儲存的路徑。

此引數在 is_memmap() 返回 False（例如，當 TensorDict 解鎖時）後即失效。

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

選擇 TensorDict 的鍵，並返回一個僅包含所選鍵的新 TensorDict。

值不會被複制：對原始tensordict或新tensordict的張量的就地修改將導致兩個tensordict都發生變化。

引數:

*keys (str) – 要選擇的鍵。
inplace (bool) – 如果為 True，則 tensordict 將就地剪枝。預設為 False。
strict (bool, optional) – 選擇一個不存在的鍵是否會返回錯誤。預設為 True。

返回:

一個新的 tensordict（如果 inplace=True 則為同一 tensordict），僅包含選定的鍵。

注意

To select keys in a tensordict and return a version of this tensordict deprived of these keys, see the split_keys() method.

示例

>>> 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: 'torch.distributed.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → None¶

將 tensordict 的內容傳送到遠端工作程序。

引數:

dst (int) – 應將內容傳送到的目標工作程序的 rank。

關鍵字引數:

group (torch.distributed.ProcessGroup, 可選) – 如果設定，將使用指定的程序組進行通訊。否則，將使用預設程序組。預設為 `None`。
init_tag (int) – 用於標記張量的初始 init_tag。注意，此值將根據 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://: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://: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()

separates(*keys: NestedKey, default: Any = _NoDefault.ZERO, strict: bool = True, filter_empty: bool = True) → Any¶

將指定的鍵從 tensordict 中分離（就地操作）。

另請參閱

此方法等效於在單個分割上呼叫 split_keys() 且 inplace=True。

另請參閱

此方法等效於呼叫 exclude()，只是它返回資料的另一個分割。

引數:

keys (NestedKey) – 要從 tensordict 中分離的鍵。
default (Any, optional) – 當鍵丟失時返回的值。如果未指定且 strict=True，則會引發異常。否則，任何缺失鍵的預設值將是 None，除非另有指定。
strict (bool, optional) – 如果為 True，則在鍵丟失時引發異常。預設為 True。
filter_empty (bool, optional) – 如果為 True，則會刪除 self 中的空 tensordicts。預設為 True。

返回:

分離出的 tensordict。

返回型別:

T

示例

>>> td = TensorDict(
...     a=0,
...     b=0,
...     c=0,
...     d=0,
... )
>>> td_a_c = td.separates("a", "c")
>>> print(td_a_c)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(td)
TensorDict(
    fields={
        b: 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)

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

設定一個新的鍵值對。

引數:

key (str, str 元組) – 要設定的鍵的名稱。
item (torch.Tensor or equivalent, TensorDictBase instance) – 要儲存在 tensordict 中的值。
inplace (bool, optional) – if True and if a key matches an existing key in the tensordict, then the update will occur in-place for that key-value pair. If inplace is True and the entry cannot be found, it will be added. For a more restrictive in-place operation, use set_() instead. Defaults to False.

關鍵字引數:

non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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: Union[Tensor, TensorCollection]) → Any¶

設定現有鍵的值，同時保留原始儲存。

引數:

key (str) – 值的名稱
item (torch.Tensor or compatible type, TensorDictBase) – 要儲存在 tensordict 中的值。

關鍵字引數:

non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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: Union[Tensor, TensorCollection], index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]]) → Any¶

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

引數:

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

關鍵字引數:

non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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 中包含的額外狀態。

This function is called from load_state_dict() to handle any extra state found within the state_dict. Implement this function and a corresponding get_extra_state() for your module if you need to store extra state within its state_dict.

引數:: state (dict) – state_dict 中的額外狀態

set_non_tensor(key: NestedKey, value: Any)¶

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

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

返回：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, strict: bool = False) → None¶

如果存在，設定由 target 給定的子模組，否則丟擲錯誤。

注意

如果 strict 設定為 False（預設），該方法將替換現有子模組或在父模組存在的情況下建立新子模組。如果 strict 設定為 True，該方法將僅嘗試替換現有子模組，並在子模組不存在時引發錯誤。

例如，假設您有一個 nn.Module A，它看起來像這樣

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(3, 3, 3)
        )
        (linear): Linear(3, 3)
    )
)

(圖示了一個 nn.Module A。 A 包含一個巢狀子模組 net_b，該子模組本身有兩個子模組 net_c 和 linear。 net_c 隨後又有一個子模組 conv。)

要用一個新的 Linear 子模組覆蓋 Conv2d，可以呼叫 set_submodule("net_b.net_c.conv", nn.Linear(1, 1))，其中 strict 可以是 True 或 False。

要將一個新的 Conv2d 子模組新增到現有的 net_b 模組中，可以呼叫 set_submodule("net_b.conv", nn.Conv2d(1, 1, 1))。

在上面，如果設定 strict=True 並呼叫 set_submodule("net_b.conv", nn.Conv2d(1, 1, 1), strict=True)，則會引發 AttributeError，因為 net_b 中不存在名為 conv 的子模組。

引數:

target – 要查詢的子模組的完全限定字串名稱。（要指定完全限定字串，請參閱上面的示例。）
module – 要設定子模組的物件。
strict – 如果為 False，該方法將替換現有子模組或建立新子模組（如果父模組存在）。如果為 True，則該方法只會嘗試替換現有子模組，如果子模組不存在則丟擲錯誤。

丟擲:

ValueError – If the target string is empty or if module is not an instance of nn.Module.
AttributeError – If at any point along the path resulting from the target string the (sub)path resolves to a non-existent attribute name or an object that is not an instance of nn.Module.

setdefault(key: NestedKey, default: Union[Tensor, TensorCollection, Any], inplace: bool = False) → Union[Tensor, TensorCollection]¶

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

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

引數:

key (str 或 nested key) – 值的名稱。
default (torch.Tensor or compatible type, TensorDictBase) – 如果鍵不存在，則儲存在 tensordict 中的值。

返回:

tensordict 中 key 的值。如果 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() → Self¶: 請參閱 torch.Tensor.share_memory_()。

share_memory_(*args, **kwargs)¶

將所有張量放入共享記憶體。

TensorDict 將被鎖定，這意味著任何非原地寫入操作都會引發異常（例如，重新命名、設定或刪除條目）。反之，一旦 tensordict 被解鎖，share_memory 屬性將變為 False，因為程序間的身份不再有保證。

返回:: self

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

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

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

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

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

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

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

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

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

返回由 dim 指定維度的尺寸。

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

softmax(dim: int, dtype: Optional[dtype] = None)¶

將 softmax 函式應用於 tensordict 元素。

引數:

dim (int 或 int 元組) – 將在其上計算 softmax 的 tensordict 維度。
dtype (torch.dtype, optional) – 返回張量的期望資料型別。如果指定，則在執行操作之前，輸入張量將被轉換為 dtype。這有助於防止資料型別溢位。

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

按字母順序返回鍵。

不支援其他引數。

如果 TensorDict 被鎖定，鍵將被快取，直到 TensorDict 解鎖以獲得更快的執行。

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

像 torch.split 一樣，使用指定的尺寸在給定維度上分割 TensorDict 中的每個張量。

返回一個包含分割塊檢視的 TensorDict 例項列表。

引數:

split_size (int or 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, default: Any = _NoDefault.ZERO, strict: bool = True, reproduce_struct: bool = False) → Tuple[T, ...]¶

根據一個或多個鍵集將 tensordict 分割成子集。

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

引數:

key_sets (鍵的序列 Dict[in_key, out_key] 或 鍵列表) – 各個分割。
inplace (bool, optional) – 如果為 True，則 self 中的鍵將被原地移除。預設為 False。
default (Any, 可選) – 當 key 丟失時要返回的值。如果未指定且 strict=True，則會引發異常。
strict (bool, optional) – 如果為 True，則在鍵丟失時引發異常。預設為 True。
reproduce_struct (bool, optional) – 如果為 True，則所有返回的 tensordict 都將具有與 self 相同的樹結構，即使某些子 tensordict 不包含葉子。

注意

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() → Any¶: 計算 self 的逐元素平方根。

sqrt_() → Any¶: In-place version of sqrt().

squeeze(*args, **kwargs)¶

壓縮維度在 -self.batch_dims+1 和 self.batch_dims-1 之間的所有張量，並將它們返回到一個新的 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 的更改將是異變的，即原始張量的內容不會被修改。這也假定 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: int = 0, *, out=None)¶

沿給定維度將 tensordicts 堆疊成一個單一的 tensordict。

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

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

將 tensordict 的所有條目堆疊成一個張量。

引數:

dim (int, optional) – 條目應堆疊的維度。

關鍵字引數:

sorted (bool or list of NestedKeys) – if True, the entries will be stacked in alphabetical order. If False (default), the dict order will be used. Alternatively, a list of key names can be provided and the tensors will be stacked accordingly. This incurs some overhead as the list of keys will be checked against the list of leaf names in the tensordict.
out (torch.Tensor, optional) – 用於堆疊操作的可選目標張量。

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

將條目堆疊到新的條目中，並可能移除原始值。

引數:

keys (NestedKey 序列) – 要堆疊的條目。

關鍵字引數:

out_key (NestedKey) – 堆疊輸入的新的鍵名。
keep_entries (bool, optional) – 如果為 False，則 keys 中的條目將被刪除。預設為 False。
dim (int, optional) – the dimension along which the stack must occur. Defaults to 0.

返回: 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。

state_dict 的結構將保持巢狀，除非將 flatten 設定為 True。

tensordict state_dict 包含重建 tensordict 所需的所有張量和元資料（目前不支援名稱）。

引數:

destination (dict, optional) – 如果提供，tensordict 的狀態將更新到 dict 中，並返回相同的物件。否則，將建立一個 OrderedDict 並返回。預設：None。
prefix (str, optional) – 新增到張量名稱的字首，用於組合 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], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回輸入 tensordict 中所有元素的標準差值。

引數:

dim (int, tuple of int, optional) – 如果為 None，則返回一個無量綱的 tensordict，其中包含所有葉子節點的總和值（如果可計算）。如果為整數或整數元組，則僅當該維度與 tensordict 形狀相容時，才會在指定的維度上呼叫 std。目前只允許 “feature” 字串。使用 dim=”feature” 將在所有特徵維度上進行歸約。如果 reduce=True，將返回一個形狀與 TensorDict 的 batch-size 相同的張量。否則，將返回一個結構與 self 相同的、具有歸約特徵維度的新 tensordict。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

correction (int) – difference between the sample size and sample degrees of freedom. Defaults to Bessel’s correction, correction=1.
reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.std(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.std()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.std(reduce=True)
tensor(1.0006)
>>> td.std(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.std(reduce=True, dim="feature")
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
>>> td.std(reduce=True, dim=0)
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])

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

將 other 縮放 alpha 後，從 self 中減去。

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

支援廣播、型別提升以及整數、浮點數和複數輸入。

引數:

other (TensorDict, Tensor 或 Number) – 從 self 中減去的張量或數字。

關鍵字引數:

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

sub_(other: tensordict.base.TensorDictBase | torch.Tensor | float, alpha: float | None = None)¶: In-place version of sub().

注意

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

sum(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回輸入 tensordict 中所有元素的總和值。

引數:

dim (int, tuple of int, optional) – 如果 None，則返回一個無維度的 tensordict，其中包含所有葉子的總和（如果可以計算）。如果為整數或整數元組，則僅當此維度與 tensordict 的形狀相容時，才會對指定的維度呼叫 sum。目前僅允許 “feature” 字串。使用 dim=”feature” 將實現跨所有特徵維度的約簡。如果 reduce=True，則將返回一個形狀與 TensorDict 的 batch_size 相同的張量。否則，將返回一個具有與 self 相同的結構但特徵維度已約簡的新 tensordict。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定，則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設：None。
reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.sum(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.sum()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.sum(reduce=True)
tensor(-0.)
>>> td.sum(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.sum(reduce=True, dim="feature")
tensor([[15., 15., 15., 15.],
        [15., 15., 15., 15.],
        [15., 15., 15., 15.]])
>>> td.sum(reduce=True, dim=0)
tensor([[9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.]])

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

tan_() → Any¶: 原地計算 TensorDict 中每個元素的 tan() 值。

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

tanh_() → Any¶: 原地計算 TensorDict 中每個元素的 tanh() 值。

tensor_split(indices_or_sections: int | list[int] | tuple[int, ...] | torch.Tensor, dim=0) → tuple[tensordict.base.TensorDictBase, ...]¶

將 TensorDict 沿著 dim 維度根據 indices_or_sections 指定的索引或節數分割成多個子 tensordict，所有這些子 tensordict 都是輸入張量的檢視。

引數:

indices_or_sections (int or List(int) or tuple(int) or 1D tensor of ints) – If indices_or_sections is an integer n or a zero dimensional long tensordict with value n, input is split into n sections along dimension dim. If input is divisible by n along dimension dim, each section will be of equal size, input.size(dim) / n. If input is not divisible by n, the sizes of the first int(input.size(dim) % n) sections will have size int(input.size(dim) / n) + 1, and the rest will have size int(input.size(dim) / n). If indices_or_sections is a list or tuple of ints, or a one-dimensional long tensor, then input is split along dimension dim at each of the indices in the list, tuple or tensor. For instance, indices_or_sections=[2, 3] and dim=0 would result in the tensors input[:2], input[2:3], and input[3:]. If indices_or_sections is a tensor, it must be a zero-dimensional or one-dimensional long tensor on the CPU.
dim (int, optional) – 要分割張量的維度。預設為 0

示例

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

to(*args, **kwargs) → TensorDictBase¶

將 TensorDictBase 子類對映到另一個裝置、dtype 或另一個 TensorDictBase 子類（如果允許）。

不允許將張量轉換為新的 dtype，因為 tensordict 不會繫結到只包含單一 tensor dtype。

引數:

device (torch.device, optional) – tensordict 的期望裝置。
dtype (torch.dtype, optional) – tensordict 的期望浮點數或複數 dtype。
tensor (torch.Tensor, optional) – 此 TensorDict 中所有張量的期望 dtype 和裝置都取自此張量。

關鍵字引數:

non_blocking (bool, optional) – whether the operations should be blocking.
memory_format (torch.memory_format, optional) – 此 tensordict 中 4D 引數和緩衝區所需的記憶體格式。
batch_size (torch.Size, optional) – 輸出 tensordict 的結果 batch-size。
other (TensorDictBase, optional) –
TensorDict 例項，其 dtype 和 device 是該 TensorDict 中所有張量的目標 dtype 和 device。

注意

由於 TensorDictBase 例項沒有 dtype，因此 dtype 是從示例葉子節點收集的。如果存在多個 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) – if non_blocking_pin=True, the number of threads to be used for pin_memory. By default, max(1, torch.get_num_threads()) threads will be spawn. num_threads=0 will cancel any multithreading for the pin_memory() calls.
inplace (bool, optional) – if True, the data will be written in-place in the same tensordict. This can be significantly faster whenever building a tensordict is CPU-overhead bound. Defaults to False.

返回:

如果裝置與 tensordict 裝置不同，以及/或傳遞了 dtype，則返回一個新的 tensordict 例項。否則，返回相同的 tensordict。batch_size 的修改是就地進行的。

注意

如果 TensorDict 是合併的，結果 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, convert_tensors: Union[bool, Literal['numpy']] = False, tolist_first: bool = False) → dict[str, Any]¶

返回一個鍵值對與 tensordict 匹配的字典。

引數:

retain_none (bool) – 如果為 True，則 tensorclass 例項中的 None 值將被寫入字典。否則，它們將被丟棄。預設值：True。
convert_tensors (bool, "numpy") – 如果為 True，則在建立字典時，張量將被轉換為列表。如果為“numpy”，則張量將被轉換為 numpy 陣列。否則，它們將保持為張量。預設值：False。
tolist_first (bool) – 如果為 True，則當 tensordict 具有批處理維度時，它將首先轉換為列表。預設值：False。

返回:

tensordict 的字典表示。

另請參閱

tolist()

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     a=torch.arange(24).view(2, 3, 4),
...     b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)),
...     batch_size=(2, 3)
... )
>>> print(td.to_dict())
{'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]]]), ‘b’: {‘c’: tensor([[[ 0, 1], [ 2, 3], [ 4, 5]],

[[ 6, 7],
[ 8, 9], [10, 11]]])}}

>>> print(td.to_dict(convert_tensors=True))
{'a': [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]], 'b': {'c': [[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]}}

to_empty(*, device: Optional[Union[int, str, device]], recurse: bool = True) → Self¶

將引數和緩衝區移動到指定裝置，而不復制儲存。

引數:

device (torch.device) – The desired device of the parameters and buffers in this module. – 此模組中引數和緩衝區的目標裝置。
recurse (bool) – 是否遞迴地將子模組的引數和緩衝區移動到指定的裝置。

返回:

self

返回型別:

模組

to_h5(filename, **kwargs)¶

將 tensordict 轉換為具有 h5 後端的 PersistentTensorDict。

引數:

filename (str 或 path) – h5 檔案的路徑。
**kwargs – 傳遞給 h5py.File.create_dataset() 的其他引數。

返回:

指向新建立檔案的 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_lazystack(dim: int = 0)¶

將 TensorDict 轉換為 LazyStackedTensorDict 或等價物。

注意

此方法可用於更改 LazyStackedTensorDict 的堆疊維度。例如，如果您有一個 stack_dim=1 的 LazyStackedTensorDict，您可以使用此方法將其更改為 stack_dim=0。

>>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3))
>>> td2 = td.to_lazystack()
>>> td2.batch_size
torch.Size([2, 3])
>>> assert isinstance(td2, LazyStackedTensorDict)
>>> assert td2.stack_dim == 0
>>> td3 = td2.to_lazystack(1)
>>> assert td3.stack_dim == 1
>>> td3.batch_size
torch.Size([2, 3])

引數:: dim (int, optional) – 要堆疊 tensordict 的維度。預設為 0。
返回:: LazyStackedTensorDict 例項。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3))
>>> td2 = td.to_lazystack()
>>> td2.batch_size
torch.Size([2, 3])
>>> assert isinstance(td2, LazyStackedTensorDict)

將 TensorCollection 的內容寫入流式資料集。

關鍵字引數:

out (str | Tuple[str, str]) –
用於儲存分片檔案的輸出資料集目錄。
1. 如果 out 是本地目錄，則分片檔案將本地儲存。
2. 如果 out 是遠端目錄，則將建立一個本地臨時目錄來
  快取分片檔案，然後將分片檔案上傳到遠端位置。完成後，在分片上傳後刪除臨時目錄。
3. 如果 out 是 (local_dir, remote_dir) 的元組，則分片檔案將儲存在
  local_dir 中，並上傳到遠端位置。
columns (dict[str, str]) – 可選的列字典。如果未提供，將自動從張量集合中推斷。
keep_local (bool) – 如果資料集已上傳，是否在上傳後保留本地資料集目錄。預設為 False。
compression (str, optional) – 可選的壓縮或壓縮級別。預設為 None。
hashes (List[str], optional) – 可選的要應用於分片檔案的雜湊演算法列表。預設為 None。
size_limit (Union[int, str], optional) – 可選的分片大小限制，超過此限制將開始新的分片。如果為 None，則將所有內容放入一個分片。也可以指定位元組的可讀格式，例如 "100kb" 表示 100 千位元組（100*1024）等等。預設為 1 << 26。如果
writer – (MDSWriter, optional): 要使用的寫入器。將從 out 關鍵字引數和其他輸入關鍵字引數建立。
**kwargs (Any) – 寫入器的其他設定。

注意

MDSWriter 對巢狀字典的支援有限。處理巢狀 tensordict 的正確方法是先使用 flatten_keys() 方法，然後在讀取後使用 unflatten_keys() 方法。

警告

此方法需要安裝 mosaicml-streaming。

警告

對於非張量資料，資料型別必須是固定的。tensordict 恢復資料型別的方式是檢視列表的第一個元素。如果為 None，將引發錯誤。否則，列表中的所有資料都必須具有相同的型別（或為 None 表示缺失）。

另請參閱

請參閱 Mosaic streaming 庫 API https://docs.mosaicml.com/projects/streaming

以下示例展示瞭如何建立資料集並在 PyTorch dataloader 中載入它的端到端示例。

示例

>>> import tempfile
>>> from typing import Any
>>> from tensordict import TensorDict, LazyStackedTensorDict
>>> import torch
>>>
>>>
>>> td = LazyStackedTensorDict(
...     TensorDict(a=0, b=1, c=torch.randn(2), d="a string"),
...     TensorDict(a=0, b=1, c=torch.randn(3), d="another string"),
...     TensorDict(a=0, b=1, c=torch.randn(3), d="yet another string"),
... )
>>>
>>> with tempfile.TemporaryDirectory() as tempdir:
...     # Create a dataset on one process / thread / node...
...     td.to_mds(out=tempdir)
...
...     # Create a dataloader
...     from streaming import StreamingDataset
...
...     # Load the dataset on another thread / process / node...
...     dataset = StreamingDataset(local=tempdir, remote=None, batch_size=2)
...
...     # Use the class `from_list` method as a collate_fn
...     dl = torch.utils.data.DataLoader(dataset=dataset, batch_size=2, collate_fn=LazyStackedTensorDict.from_list)
...     for batch in dl:
...         print("batch", batch)

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) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

示例

>>> 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()

返回:: 一個 tensordict，如果 return_swap 為 True，則包含來自模組的值，否則為 None。

to_namedtuple(dest_cls: Optional[type] = None) → Any¶

將 tensordict 轉換為命名元組。

引數:: dest_cls (Type, optional) – 一個可選的命名元組類。

示例

>>> 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) → Any¶

將所有巢狀張量轉換為填充版本並相應地調整批次大小。

引數:

padding (float) – tensordict 中張量的填充值。預設為 0.0。
mask_key (NestedKey, optional) – 如果提供，則為寫入有效值掩碼的鍵。如果異構維度不是 tensordict batch-size 的一部分，則會導致錯誤。預設為 None

to_pytree()¶

將 tensordict 轉換為 PyTree。

如果 tensordict 不是從 pytree 建立的，此方法將僅返回 self 而不進行修改。

有關更多資訊和示例，請參閱 from_pytree()。

to_struct_array() → ndarray¶

將 tensordict 轉換為 numpy 結構化陣列。

在 from_struct_array() - to_struct_array() 迴圈中，輸入和輸出陣列的內容應匹配。但是，to_struct_array 不會保留原始陣列的記憶體內容。

另請參閱

有關更多資訊，請參閱 from_struct_array()。

另請參閱

轉換為 numpy 陣列字典，請參閱 numpy()。

返回:: 輸入 TensorDict 的 numpy 結構化陣列表示。

示例

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({'a': torch.tensor([1, 2, 3]), 'b': torch.tensor([4.0, 5.0, 6.0])}, batch_size=[3])
>>> arr = td.to_struct_array()
>>> print(arr)
[(1, 4.) (2, 5.) (3, 6.)]

to_tensordict(*, retain_none: bool | None = None)¶

從 TensorDictBase 返回一個常規的 TensorDict 例項。

引數:: retain_none (bool) – 如果為 True，則 tensorclass 例項中的 None 值將被寫入 tensordict。否則，它們將被丟棄。預設值：True。
返回:: 包含相同值的新 TensorDict 物件。

tolist(*, convert_nodes: bool = True, convert_tensors: Union[bool, Literal['numpy']] = False, tolist_first: bool = False, as_linked_list: bool = False) → List[Any]¶

返回 tensordict 的巢狀列表表示。

如果 tensordict 沒有批次維度，則此方法返回單個列表或字典。否則，它返回一個巢狀列表，其中每個內部列表代表一個批次維度。

引數:

convert_nodes (bool) – 如果為 True，則葉子節點將被轉換為字典。否則，它們將被作為值列表返回。預設值：True。
convert_tensors (bool, "numpy") – 如果為 True，則在建立字典時，張量將被轉換為列表。如果為“numpy”，則張量將被轉換為 numpy 陣列。否則，它們將保持為張量。預設值：False。
tolist_first (bool) – 如果為 True，則當 tensordict 具有批處理維度時，它將首先轉換為列表。預設值：False。
as_linked_list (bool) – 如果為 True，則列表將被轉換為 tensordict.utils.LinkedList，它將在修改列表時自動更新 tensordict。預設值：False。

返回:

tensordict 的巢狀列表表示。

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     a=torch.arange(24).view(2, 3, 4),
...     b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)),
...     batch_size=(2, 3)
... )
>>> print(td.tolist(tolist_first=True))
[[{'a': tensor([0, 1, 2, 3]), 'b': [{'c': tensor(0)}, {'c': tensor(1)}]}, {'a': tensor([4, 5, 6, 7]), 'b': [{'c': tensor(2)}, {'c': tensor(3)}]}, {'a': tensor([ 8,  9, 10, 11]), 'b': [{'c': tensor(4)}, {'c': tensor(5)}]}], [{'a': tensor([12, 13, 14, 15]), 'b': [{'c': tensor(6)}, {'c': tensor(7)}]}, {'a': tensor([16, 17, 18, 19]), 'b': [{'c': tensor(8)}, {'c': tensor(9)}]}, {'a': tensor([20, 21, 22, 23]), 'b': [{'c': tensor(10)}, {'c': tensor(11)}]}]]
>>> print(td.tolist(tolist_first=False))
[[{'a': tensor([0, 1, 2, 3]), 'b': {'c': tensor([0, 1])}}, {'a': tensor([4, 5, 6, 7]), 'b': {'c': tensor([2, 3])}}, {'a': tensor([ 8,  9, 10, 11]), 'b': {'c': tensor([4, 5])}}], [{'a': tensor([12, 13, 14, 15]), 'b': {'c': tensor([6, 7])}}, {'a': tensor([16, 17, 18, 19]), 'b': {'c': tensor([8, 9])}}, {'a': tensor([20, 21, 22, 23]), 'b': {'c': tensor([10, 11])}}]]
>>> print(td.tolist(convert_tensors=False))
[[{'a': [0, 1, 2, 3], 'b': [{'c': 0}, {'c': 1}]}, {'a': [4, 5, 6, 7], 'b': [{'c': 2}, {'c': 3}]}, {'a': [8, 9, 10, 11], 'b': [{'c': 4}, {'c': 5}]}], [{'a': [12, 13, 14, 15], 'b': [{'c': 6}, {'c': 7}]}, {'a': [16, 17, 18, 19], 'b': [{'c': 8}, {'c': 9}]}, {'a': [20, 21, 22, 23], 'b': [{'c': 10}, {'c': 11}]}]]
>>> print(td.tolist(convert_nodes=False))
[[[tensor([0, 1, 2, 3]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([4, 5, 6, 7]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([ 8,  9, 10, 11]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)]], [[tensor([12, 13, 14, 15]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([16, 17, 18, 19]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([20, 21, 22, 23]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)]]]

train(mode: bool = True) → Self¶

將模組設定為訓練模式。

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc. – 這隻對某些模組有影響。有關其在訓練/評估模式下的行為的詳細資訊，例如它們是否受影響，請參閱特定模組的文件，例如 Dropout、BatchNorm 等。

引數:: mode (bool) – 設定訓練模式（True）或評估模式（False）。預設：True。
返回:: self
返回型別:: 模組

transpose(dim0, dim1)¶

返回輸入張量的轉置張量。給定的維度 dim0 和 dim1 將被交換。

轉置 tensordict 的原地或非原地修改也會影響原始 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() → Any¶: 計算 TensorDict 每個元素的 trunc() 值。

trunc_() → Any¶: 原地計算 TensorDict 每個元素的 trunc() 值。

type(dst_type: dtype) → Any¶

將所有張量轉換為 dst_type。

引數:: dst_type (type or string) – 所需型別

uint16() → Any¶: 將所有張量轉換為 torch.uint16。

uint32() → Any¶: 將所有張量轉換為 torch.uint32。

uint64() → Any¶: 將所有張量轉換為 torch.uint64。

uint8() → Any¶: 將所有張量轉換為 torch.uint8。

unbind(dim: int) → tuple[T, ...]¶

返回一個已解綁的 tensordict 索引元組，沿指定的維度。

示例

>>> 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) – 指定要展平的輸入張量的維度。
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 例項。巢狀 tensordicts 的元資料將從根推斷：資料樹中的所有例項將共享相同的批次大小、維度名稱和裝置。

引數:

separator (str, optional) – 巢狀項之間的分隔符。
inplace (bool, optional) – 如果為 True，則結果 tensordict 將與呼叫它的 tensordict 具有相同的身份。預設為 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_() → Any¶

解鎖 tensordict 以進行非原地操作。

可用作裝飾器。

有關更多詳細資訊，請參閱 lock_()。

unsqueeze(*args, **kwargs)¶

為介於 -td.batch_dims 和 td.batch_dims 之間的維度擴充套件所有張量，並將它們返回到一個新的 tensordict 中。

引數:: dim (int) – 要擴充套件的維度

示例

>>> 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])

此操作也可以用作上下文管理器。對原始 tensordict 的更改將是就地進行的，即原始張量的內容不會被修改。這也假定 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, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]] | 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, update_batch_size: bool = False, ignore_lock: bool = False) → TensorDictBase¶

使用來自字典或其他 TensorDict 的值來更新 TensorDict。

警告

update 在 try/except 塊中呼叫時會損壞資料。不要在這些塊中使用此方法來嘗試捕獲和修補執行過程中發生的錯誤。

引數:

input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。
clone (bool, 可選) – 在設定輸入（ tensor）字典中的張量之前是否克隆它們。預設為 False。
inplace (bool, 可選) – 如果為 True 並且某個鍵與 tensordict 中的現有鍵匹配，則更新將對該鍵值對進行原地更新。如果找不到該條目，則會將其新增。預設為 False。

關鍵字引數:

keys_to_update (NestedKeys 序列, 可選) – 如果提供，則僅更新 key_to_update 中的鍵列表。這旨在避免呼叫 data_dest.update(*keys_to_update))。
non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.
is_leaf (Callable[[Type], bool], optional) –
一個指示物件型別是應被視為葉子節點並被交換，還是應被視為張量集合的可呼叫物件。

另請參閱

is_leaf_nontensor() 和 default_is_leaf()。
update_batch_size (bool, optional) –
如果為 True，則 update 將嘗試更新目標（self）的 batch_size，如果它與源的 batch_size 不匹配。預設為 False。

注意

在 batch_size 不匹配的情況下，LazyStackTensorDict 例項將被清空其內容，並將使用源 tensordict 的副本重新填充容器。

注意

此引數假定 keys_to_update 為空，並且 inplace=False。如果目標（self）的鍵不是源鍵的子集，則會引發異常，因為 TensorDict 將無法推斷如何處理額外的目標條目。
ignore_lock (bool, optional) – 如果為 True，則可以更新任何 tensordict，無論其鎖定狀態如何。預設為 False。

注意

當使用 N 個元素的 LazyStackedTensorDict 更新另一個 M 個元素的 LazyStackedTensorDict 時，沿堆疊維度，update 方法會將額外 tensordict 的副本附加到目標（self）的懶惰堆疊中。這使得使用者可以依賴 update 來逐步增加懶惰堆疊。

返回:: 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, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], 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, 可選) – 在設定輸入（ tensor）字典中的張量之前是否克隆它們。預設為 False。

關鍵字引數:

keys_to_update (NestedKeys 序列, 可選) – 如果提供，則僅更新 key_to_update 中的鍵列表。這旨在避免呼叫 data_dest.update_(data_src.select(*keys_to_update))。
non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], 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, 可迭代物件, slice) – 更新應發生的 tensordict 的索引。
clone (bool, 可選) – 在設定輸入（ tensor）字典中的張量之前是否克隆它們。預設為 False。

關鍵字引數:

keys_to_update (NestedKeys序列, 可選) – 如果提供，將僅更新 key_to_update 中的鍵列表。
non_blocking (bool, optional) – if True and this copy is between different devices, the copy may occur asynchronously with respect to the host.

返回:

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[Union[Tensor, TensorCollection]]¶

返回一個表示 tensordict 值的生成器。

引數:

include_nested (bool, 可選) – 如果為 `True`，則返回巢狀值。預設為 `False`。
leaves_only (bool, 可選) – 如果為 `False`，則僅返回葉子節點。預設為 `False`。
is_leaf (callable, optional) –
一個作用於類型別的可呼叫物件，返回一個布林值，指示該類是否應被視為葉子節點。

注意

is_leaf 的目的是不是阻止遞迴呼叫巢狀的 tensordicts，而是為過濾標記某些型別，當 `leaves_only=True` 時。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，巢狀 tensordict 的結構仍將被遍歷。換句話說，`is_leaf` 不控制遞迴深度，而是提供了一種當 `leaves_only=True` 時從結果中過濾掉某些型別的方法。這意味著樹中的一個節點既可以是葉節點，也可以是具有子節點的節點。實際上，`is_leaf` 的預設值不包含 tensordict 和 tensorclass 例項作為葉節點。

另請參閱

is_leaf_nontensor() 和 default_is_leaf()。

關鍵字引數:

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

var(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) → Union[Any, Tensor]¶

返回輸入 tensordict 中所有元素的方差值。

引數:

dim (int, int 元組, 可選) – 如果為 None，則返回一個無量綱的 tensordict，其中包含所有葉子的總和值（如果可以計算）。如果為整數或整數元組，則當此維度與 tensordict 形狀相容時，將在指定的維度上呼叫 var。當前只允許 “feature” 字串。使用 dim=”feature” 將在所有特徵維度上執行歸約。如果 reduce=True，則將返回形狀與 TensorDict 的 batch_size 相同的張量。否則，將返回一個結構與 self 相同的、特徵維度已歸約的新 tensordict。
keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

correction (int) – difference between the sample size and sample degrees of freedom. Defaults to Bessel’s correction, correction=1.
reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上執行歸約，並返回一個單一的歸約張量。預設為 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.var(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.var()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.var(reduce=True)
tensor(1.0006)
>>> td.var(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.var(reduce=True, dim="feature")
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
>>> td.var(reduce=True, dim=0)
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])

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 – 可迭代物件

關鍵字引數:

batch_size (torch.Size, 可選) – 如果提供了 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: Tensor, other: torch.Tensor | tensordict.base.TensorDictBase, *, out: tensordict.base.TensorDictBase | None = None, pad: int | bool = None, update_batch_size: bool = False)¶

根據 condition 從 self 或 other 中選擇元素，並返回一個 TensorDict。

引數:

condition (BoolTensor) – 當 True (非零) 時，返回 self，否則返回 other。
other (TensorDictBase 或標量) – 當 other 是標量時，值為標量；或者，當 condition 為 False 時的值為選定的值。

關鍵字引數:

out (TensorDictBase, 可選) – 輸出 TensorDictBase 例項。
pad (標量, 可選) – 如果提供，源或目標 tensordict 中缺失的鍵將寫入為 torch.where(mask, self, pad) 或 torch.where(mask, pad, other)。預設為 None，即不允許缺失的鍵。
update_batch_size (bool, optional) – 如果為 True 且提供了 out，則輸出的批次大小將更新為匹配條件的批次大小。預設為 False。

xpu(device: Optional[Union[device, int]] = None) → Self¶

將所有模型引數和緩衝區移動到 XPU。

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

注意

此方法就地修改模組。

引數:: device (int, optional) – 如果指定，所有引數都將複製到該裝置
返回:: self
返回型別:: 模組

zero_() → Any¶: 原地將 tensordict 中的所有張量清零。

zero_grad(set_to_none: bool = True) → Any¶

遞迴地將 TensorDict 的所有梯度清零。

引數:: set_to_none (bool, 可選) – 如果為 True，則 tensor.grad 將為 None，否則為 0。預設為 True。

TensorDictParams¶

文件

教程

資源