PersistentTensorDict

class tensordict.PersistentTensorDict(*, batch_size=None, filename=None, group=None, mode='r', backend='h5', device=None, **kwargs)

持久化 TensorDict 實現。

PersistentTensorDict 例項提供了一個與儲存在磁碟上的資料互動的介面，使得對這些資料的訪問變得容易，同時仍然利用了後端提供的快速訪問。

與其他 TensorDictBase 子類一樣，PersistentTensorDict 具有一個 device 屬性。這不意味著資料儲存在該裝置上，而是意味著載入時，資料將被轉換為所需的裝置。

關鍵字引數:

• batch_size (torch.Size 或 相容) – tensordict 的批大小。預設為 torch.Size(())。

• filename (str, 可選) – h5 檔案的路徑。與 group 互斥。

• group (h5py.Group, 可選) – 包含資料的檔案或組。與 filename 互斥。

• mode (str, 可選) – 讀取模式。預設為 "r"。

• backend (str, 可選) – 儲存後端。目前只支援 "h5"。

• device (torch.device 或 相容, 可選) – tensordict 的裝置。預設為 None（即預設 PyTorch 裝置）。

• **kwargs – 要傳遞給 h5py.File.create_dataset() 的 kwargs。

注意

目前，PersistentTensorDict 例項在超出範圍時不會關閉。這意味著使用者有責任在必要時關閉它們。

示例

>>> import tempfile
>>> with tempfile.NamedTemporaryFile() as f:
...     data = PersistentTensorDict(file=f, batch_size=[3], mode="w")
...     data["a", "b"] = torch.randn(3, 4)
...     print(data)

abs() → T

計算 TensorDict 中每個元素的絕對值。

abs_() → T

就地計算 TensorDict 中每個元素的絕對值。

acos() → T

計算 TensorDict 中每個元素的 acos() 值。

acos_() → T

就地計算 TensorDict 中每個元素的 acos() 值。

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

將 other 乘以 alpha 後加到 self。

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

引數:

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

關鍵字引數:

• alpha (Number, 可選) – other 的乘數。

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

add_(other: tensordict.base.TensorDictBase | torch.Tensor | float, *, alpha: Optional[float] = None)

add() 的原地（in-place）版本。

注意

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

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

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

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

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

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

引數:

• other1 (TensorDict 或 Tensor) – 分子 tensordict (或 tensor)

• tensor2 (TensorDict 或 Tensor) – 分母 tensordict (或 tensor)

關鍵字引數:

value (Number, 可選) – \(\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)

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

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

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

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

引數:

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

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

關鍵字引數:

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

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

addcmul() 的原地版本。

all(dim: Optional[int] = None) → bool | tensordict.base.TensorDictBase

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

引數:

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

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

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

與 max() 相同，區別在於 return_indices=False。

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

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

與 min() 相同，區別在於 return_indices=False。

any(dim: Optional[int] = None) → bool | tensordict.base.TensorDictBase

檢查 tensordict 中是否存在任何 True/非空值。

引數:

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

apply(fn: Callable, *others: T, 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: Optional[bool] = None, propagate_lock: bool = False, call_on_nested: bool = False, out: Optional[TensorDictBase] = None, **constructor_kwargs) → Optional[T]

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

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

引數:

• fn (Callable) – 應用於 tensordict 中張量的函式。

• *others (TensorDictBase 例項, 可選) – 如果提供，這些 tensordict 例項應具有與 self 匹配的結構。 fn 引數應接收與 tensordict 數量一樣多的未命名輸入，包括 self。如果其他 tensordict 缺少條目，則可以透過 default 關鍵字引數傳遞預設值。

關鍵字引數:

• batch_size (int 序列, 可選) – 如果提供，生成的 TensorDict 將具有所需的 batch_size。 batch_size 引數應與轉換後的 batch_size 匹配。這是一個僅關鍵字引數。

• device (torch.device, 可選) – 生成的裝置，如果存在。

• names (str 列表, 可選) – 新的維度名稱，在 batch_size 被修改的情況下。

• inplace (bool, 可選) – 如果為 True，則進行就地更改。 預設為 False。 這是一個僅關鍵字引數。

• default (Any, 可選) – 其他 tensordict 中缺少條目的預設值。 如果未提供，缺少條目將引發 KeyError。

• filter_empty (bool, 可選) – 如果 True，則將過濾掉空的 tensordict。 這也降低了計算成本，因為不會建立和銷燬空的資料結構。 非張量資料被視為葉節點，因此即使函式未觸及，也會保留在 tensordict 中。 為了向後相容，預設為 False。

• propagate_lock (bool, 可選) – 如果 True，則鎖定的 tensordict 將生成另一個鎖定的 tensordict。 預設為 False。

• call_on_nested (bool, 可選) –

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

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

• out (TensorDictBase, 可選) –

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

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

  警告

  如果對 tensordict 執行的操作需要訪問多個鍵才能進行單次計算，則提供等於 self 的 out 引數可能會導致該操作默默地產生錯誤的結果。 例如

  >>> td = TensorDict({"a": 1, "b": 1})
  >>> td.apply(lambda x: x+td["a"])["b"] # Right!
  tensor(2)
  >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong!
  tensor(3)
  

• **constructor_kwargs – 要傳遞給 TensorDict 建構函式的其他關鍵字引數。

返回值:

具有 transformed_in 張量的新 tensordict。

示例

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

注意

如果函式返回 None，則忽略該條目。 這可用於過濾 tensordict 中的資料

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

注意

apply 方法將返回一個 TensorDict 例項，無論輸入型別如何。 為了保持相同的型別，可以執行

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

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

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

引數:

• fn (Callable) – 應用於 tensordict 中張量的函式。

• *others (TensorDictBase 序列, 可選) – 要使用的其他 tensordict。

關鍵字引數：參見 apply()。

返回值:

應用了該函式的 self 或 self 的副本

asin() → T

計算 TensorDict 中每個元素的 asin() 值。

asin_() → T

就地計算 TensorDict 中每個元素的 asin() 值。

atan() → T

計算 TensorDict 中每個元素的 atan() 值。

atan_() → T

就地計算 TensorDict 中每個元素的 atan() 值。

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

設定 tensordict 的最大批大小，最多到可選的 batch_dims。

引數:

batch_dims (int, 可選) – 如果提供，則批大小最多為 batch_dims 長。

返回值:

self

示例

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

auto_device_() → T

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

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

property batch_dims: int

tensordict 批大小的長度。

返回值:

int，描述 tensordict 的維度數量。

property batch_size

TensorDict 的形狀（或批大小）。

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

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

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

返回值:

描述 TensorDict 批大小的 Size 物件。

示例

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

bfloat16()

將所有張量轉換為 torch.bfloat16。

bitwise_and(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) → TensorDictBase

執行 self 和 other 之間的按位與操作。

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

引數:

other (TensorDictBase 或 torch.Tensor) – 用於執行按位與操作的張量或 TensorDict。

關鍵字引數:

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

bool()

將所有張量轉換為 torch.bool。

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

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

關鍵字引數:

count_duplicates (bool) – 是否將重複的張量視為獨立的。 如果 False，則只會丟棄嚴格相同的張量（來自公共基本張量的相同檢視但不同的 ID 將被計算兩次）。 預設為 True（假設每個張量都是單個副本）。

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

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

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

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

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

引數:

dim (int, optional) – 條目應沿其連線的維度。

關鍵字引數:

• sorted (bool 或 NestedKeys 列表) – 如果 True，則條目將按字母順序連線。 如果 False（預設），將使用 dict 順序。 或者，可以提供鍵名列表，並將相應地連線張量。 這會產生一些開銷，因為將根據 tensordict 中的葉名稱列表檢查鍵列表。

• out (torch.Tensor, optional) – cat 操作的可選目標張量。

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

將條目連線到新條目中，並可能刪除原始值。

引數:

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

關鍵詞引數

out_key (NestedKey): 連線輸入的新鍵名。 keep_entries (bool, optional): 如果 False，則將刪除 keys 中的條目。

預設為 False。

dim (int, optional): 必須發生連線的維度。

預設為 0。

返回: self

示例

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

ceil() → T

計算 TensorDict 的每個元素的 ceil() 值。

ceil_() → T

就地計算 TensorDict 的每個元素的 ceil() 值。

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

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

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

引數:

• chunks (int) – 要返回的塊的數量

• dim (int, optional) – 沿著該維度分割 tensordict。預設為 0。

示例

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

clamp(min: tensordict.base.TensorDictBase | torch.Tensor = None, max: tensordict.base.TensorDictBase | torch.Tensor = None, *, out=None)

將 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: str | torch.Tensor | None = None) → T

如果 self 的元素大於該值，則將其鉗制為 other。

引數:

other (TensorDict or Tensor) – 另一個輸入 tensordict 或 tensor。

關鍵字引數:

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

clamp_max_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

clamp_max() 的原地版本。

注意

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

clamp_min(other: tensordict.base.TensorDictBase | torch.Tensor, default: str | torch.Tensor | None = None) → T

如果 self 的元素小於該值，則將其鉗制為 other。

引數:

other (TensorDict or Tensor) – 另一個輸入 tensordict 或 tensor。

關鍵字引數:

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

clamp_min_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

clamp_min() 的原地版本。

注意

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

clear() → T

擦除 tensordict 的內容。

clear_device_() → T

清除 tensordict 的裝置。

返回: self

clear_refs_for_compile_() → T

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

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

返回: self

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

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

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

引數:

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

注意

與許多其他操作（逐點算術、形狀操作等）不同，clone 不繼承原始的鎖定屬性。 做出此設計選擇是為了可以建立克隆進行修改，這是最常見的用法。

close()

關閉持久化 tensordict。

complex128()

將所有張量轉換為 torch.complex128 型別。

complex32()

將所有張量轉換為 torch.complex32 型別。

complex64()

將所有張量轉換為 torch.complex64 型別。

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

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

引數:

filename (Path, optional) – 可選的檔案路徑，用於將記憶體對映張量用作 tensordict 的儲存。

關鍵字引數:

• num_threads (integer, optional) – 用於填充儲存的執行緒數。

• device (torch.device, optional) – 可選的裝置，儲存必須在該裝置上例項化。

• non_blocking (bool, optional) – 傳遞給 copy_() 的 non_blocking 引數。

• inplace (bool, optional) – 如果為 True，則生成的 tensordict 與 self 相同，並更新了值。預設為 False。

• return_early (bool, optional) – 如果為 True 且 num_threads>0，則該方法將返回 tensordict 的 future。可以使用 future.result() 查詢結果 tensordict。

• use_buffer (bool, optional) – 如果為 True 且傳遞了檔名，則將在共享記憶體中建立一箇中間本地緩衝區，並且資料將在最後一步複製到儲存位置。這可能比直接寫入遠端物理記憶體（例如 NFS）更快。預設為 False。

• share_memory (bool, optional) – 如果為 True，則儲存將被放置在共享記憶體中。預設為 False。

• pin_memory (bool, optional) – 是否應將整合後的資料放置在 pinned 記憶體中。預設為 False。

• metadata (bool, optional) – 如果為 True，則元資料將與公共儲存一起儲存。如果提供了檔名，則這無效。當人們想要控制如何實現序列化時，儲存元資料可能很有用，因為如果元資料可用或不可用，TensorDict 會以不同的方式處理整合的 TD 的 pickle/unpickle。

注意

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

在常規 TensorDict 上例項化 PersistentTensorDict。

copy()

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

等效於 TensorDictBase.clone(recurse=False)

copy_(tensordict: T, non_blocking: bool = False) → 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) → T

請參閱 TensorDictBase.update_at_。

cos() → T

計算 TensorDict 中每個元素的 cos() 值。

cos_() → T

就地計算 TensorDict 中每個元素的 cos() 值。

cosh() → T

計算 TensorDict 中每個元素的 cosh() 值。

cosh_() → T

就地計算 TensorDict 中每個元素的 cosh() 值。

cpu(**kwargs) → T

將 tensordict 轉換為 CPU。

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

create_nested(key)

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

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

示例

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

cuda(device: Optional[int] = None, **kwargs) → T

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

引數:

device (int, optional) – 如果提供，則為張量應投射到的 cuda 裝置。

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

cummax(dim: int, *, reduce: Optional[bool] = None, return_indices: bool = True) → tensordict.base.TensorDictBase | torch.Tensor

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

引數:

dim (int) – 一個整數，表示執行 cummax 操作的維度。

關鍵字引數:

• reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上進行縮減，並返回單個縮減張量。預設為 False。

• return_argmins (bool, optional) – 當傳遞 dim 引數時，cummax() 返回一個帶有值和索引的命名元組。 此的 TensorDict 等效項是返回一個張量類，其中包含具有 idendical 結構的條目 "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) → tensordict.base.TensorDictBase | torch.Tensor

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

引數:

dim (int) – 一個整數，表示執行 cummin 操作的維度。

關鍵字引數:

• reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上進行縮減，並返回單個縮減張量。預設為 False。

• return_argmins (bool, optional) – cummin() 在傳遞 dim 引數時，會返回一個包含 values 和 indices 的命名元組。 TensorDict 的等效操作是返回一個 tensorclass，其中包含 "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

返回一個 tensordict，其中包含葉張量的 .data 屬性。

data_ptr(*, storage: bool = False)

返回 tensordict 葉子的 data_ptr。

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

關鍵字引數:

storage (bool, optional) – 如果 True，則改為呼叫 tensor.untyped_storage().data_ptr()。預設為 False。

示例

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

注意

LazyStackedTensorDict 例項將顯示為巢狀的 tensordict，以反映其葉子的真實 data_ptr()

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

del_(key)

刪除 tensordict 的鍵。

引數:

key (NestedKey) – 要刪除的鍵

返回值:

self

densify(layout: layout = torch.strided)

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

關鍵字引數:

layout (torch.layout) – 巢狀張量的佈局（如果存在）。 預設為 strided。

property depth: int

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

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

detach() → T

分離 tensordict 中的張量。

返回值:

一個新的 tensordict，沒有需要梯度的張量。

detach_()

就地分離 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: str | torch.Tensor | None = None) → T

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

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

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

引數:

other (TensorDict, Tensor or Number) – 除數。

關鍵字引數:

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

div_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

div() 的就地版本。

注意

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

double()

將所有張量轉換為 torch.bool。

property dtype

返回 tensordict 中值的 dtype（如果它是唯一的）。

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

將 tensordict 儲存到磁碟。

此函式是 memmap() 的代理。

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

返回一個新的、空的 tensordict，具有相同的裝置和批次大小。

引數:

recurse (bool, optional) – 如果為 True，則將複製 TensorDict 的整個結構，而無需內容。 否則，只會複製根。預設為 False。

關鍵字引數:

• batch_size (torch.Size, optional) – tensordict 的新批次大小。

• device (torch.device, optional) – 新裝置。

• names (list of str, optional) – 維度名稱。

entry_class(key: NestedKey) → type

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

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

erf() → T

計算 TensorDict 中每個元素的 erf() 值。

erf_() → T

就地計算 TensorDict 中每個元素的 erf() 值。

erfc() → T

計算 TensorDict 中每個元素的 erfc() 值。

erfc_() → T

就地計算 TensorDict 中每個元素的 erfc() 值。

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

排除 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() → T

計算 TensorDict 中每個元素的 exp() 值。

exp_() → T

就地計算 TensorDict 中每個元素的 exp() 值。

expand(*args, **kwargs) → T

根據 expand() 函式擴充套件 tensordict 的每個張量，忽略特徵維度。

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

示例

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

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

將 tensordict 的形狀廣播到 other 的形狀，並據此進行擴充套件。

如果輸入是張量集合（tensordict 或 tensorclass），則葉子將以一對一的方式擴充套件。

示例

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

expm1() → T

計算 TensorDict 中每個元素的 expm1() 值。

expm1_() → T

就地計算 TensorDict 中每個元素的 expm1() 值。

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

使用給定的值填充由鍵指向的張量。

引數:

• key (str) – 要重新命名的鍵

• value (Number, bool) – 用於填充的值

返回值:

self

filter_empty_()

就地過濾掉所有空的tensordict。

filter_non_tensor_data() → T

過濾掉所有非張量資料。

flatten(start_dim=0, end_dim=- 1)

扁平化 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) → T

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

TensorDict 型別將會丟失，結果將是一個簡單的 TensorDict 例項。

引數:

• separator (str, 可選) – 巢狀項之間的分隔符。

• inplace (bool, 可選) – 如果 True，則結果 tensordict 將與進行呼叫的 tensordict 具有相同的標識。預設為 False。

• is_leaf (callable, optional) –

  一個可呼叫物件，作用於一個類型別，返回一個布林值，指示是否應將該類視為葉子節點。

  注意

  is_leaf 的目的不是阻止遞迴呼叫到巢狀的 tensordict 中，而是當 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() 在處理狀態字典時特別有用，因為它們可以無縫地將扁平的字典轉換為模仿模型結構的資料結構。

示例

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

將所有張量轉換為 torch.float。

float16()

將所有張量轉換為 torch.float16。

float32()

將所有張量轉換為 torch.float32。

float64()

將所有張量轉換為 torch.float64。

floor() → T

計算 TensorDict 中每個元素的 floor() 值。

floor_() → T

就地計算 TensorDict 中每個元素的 floor() 值。

frac() → T

計算 TensorDict 中每個元素的 frac() 值。

frac_() → T

就地計算 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)

遞迴地將任何物件轉換為 TensorDict。

注意

from_any 比常規的 TensorDict 建構函式限制更少。 它可以利用自定義啟發式方法將資料類或元組之類的資料結構轉換為 tensordict。 這種方法可能會產生一些額外的開銷，並且涉及更多關於對映策略的有主見的選項。

注意

此方法遞迴地將輸入物件轉換為 TensorDict。 如果該物件已經是 TensorDict（或任何類似的張量集合物件），它將按原樣返回。

引數:

obj – 要轉換的物件。

關鍵字引數:

• auto_batch_size (bool, 可選) – 如果 True，將自動計算批處理大小。 預設為 False。

• batch_dims (int, 可選) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。 預設為 None（每個級別的完整批處理大小）。

• device (torch.device, 可選) – 將在上面建立 TensorDict 的裝置。

• batch_size (torch.Size, 可選) – TensorDict 的批處理大小。 與 auto_batch_size 互斥。

返回值:

輸入物件的 TensorDict 表示。

支援的物件

• 透過 from_dataclass() 的資料類（資料類將被轉換為 TensorDict 例項，而不是 tensorclass）。

• 透過 from_namedtuple() 從命名元組建立。

• 透過 from_dict() 從字典建立。

• 透過 from_tuple() 從元組建立。

• 透過 from_struct_array() 從 NumPy 的結構化陣列建立。

• 透過 from_h5() 從 HDF5 物件建立。

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

將資料類轉換為 TensorDict 例項。

引數:

dataclass – 要轉換的資料類例項。

關鍵字引數:

• auto_batch_size (bool, 可選) – 如果為 True，則自動確定並將批處理大小應用於生成的 TensorDict。預設為 False。

• batch_dims (int, 可選) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每級完全批處理大小）。

• as_tensorclass (bool, 可選) – 如果為 True，則將轉換委託給自由函式 from_dataclass()，並返回一個與張量相容的類（tensorclass()）或例項，而不是 TensorDict。預設為 False。

• device (torch.device, 可選) – 將在其上建立 TensorDict 的裝置。預設為 None。

• batch_size (torch.Size, 可選) – TensorDict 的批處理大小。預設為 None。

返回值:

從提供的資料類派生的 TensorDict 例項，除非 as_tensorclass 為 True，在這種情況下，將返回與張量相容的類或例項。

引發:

TypeError – 如果提供的輸入不是資料類例項。

警告

此方法不同於自由函式 from_dataclass，並且用途不同。雖然自由函式返回與張量相容的類或例項，但此方法返回 TensorDict 例項。

classmethod from_dict(input_dict, filename, *others, auto_batch_size: bool = False, batch_size=None, device=None, **kwargs)

將字典或 TensorDict 轉換為 h5 檔案。

引數:

• input_dict (dict, TensorDict 或 相容型別) – 要儲存為 h5 的資料。

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

關鍵字引數:

• auto_batch_size (bool, 可選) – 如果 True，將自動計算批處理大小。 預設為 False。

• batch_size (tensordict 批處理大小, 可選) – 如果提供，則為 tensordict 的批處理大小。 如果沒有，批處理大小將從輸入結構（如果存在）中收集或自動確定。

• device (torch.device 或 相容型別, 可選) – 希望在返回張量後張量所在的裝置。預設為 None（預設情況下在 cpu 上）。

• **kwargs – 要傳遞給 h5py.File.create_dataset() 的 kwargs。

返回值:

一個連結到新建立的檔案的 PersitentTensorDict 例項。

from_dict_instance(input_dict, *others, auto_batch_size: Optional[bool] = None, batch_size=None, device=None, batch_dims=None, names=None)

from_dict() 的例項方法版本。

與 from_dict() 不同，此方法將嘗試保持現有樹中的 tensordict 型別（對於任何現有葉子節點）。

示例

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

classmethod from_h5(filename, *, mode='r', batch_size: torch.size | None = None)

從 h5 檔案建立 PersistentTensorDict。

此函式將自動確定每個巢狀 tensordict 的批處理大小（除非提供了 batch_size）。

引數:

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

關鍵字引數:

• mode (str, 可選) – 讀取模式。預設為 "r"。

• batch_size (torch.Size, 可選) – TensorDict 的批處理大小。預設為 None（自動確定批處理大小）。

返回值:

輸入 h5 檔案的 PersistentTensorDict 表示。

示例

>>> ptd = PersistentTensorDict.from_h5("path/to/file.h5")
>>> print(ptd)
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, 可選) – 如果為 True，將返回 TensorDictParams 例項，該例項可用於在 torch.nn.Module 中儲存引數。預設為 False。

• lock (bool, 可選) – 如果為 True，則生成的 tensordict 將被鎖定。預設為 True。

• use_state_dict (bool, 可選) –

  如果為 True，將使用模組的 state-dict，並將其解壓為 TensorDict，其具有模型的樹形結構。預設為 False。

  注意

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

示例

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

檢索多個模組的引數，用於 ensebmle 學習/期望應用程式的 vmap 特徵。

引數:

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

關鍵字引數:

• as_module (bool, 可選) – 如果為 True，將返回 TensorDictParams 例項，該例項可用於在 torch.nn.Module 中儲存引數。預設為 False。

• lock (bool, 可選) – 如果為 True，則生成的 tensordict 將被鎖定。預設為 True。

• use_state_dict (bool, 可選) –

  如果為 True，將使用模組的 state-dict，並將其解壓為 TensorDict，其具有模型的樹形結構。預設為 False。

  注意

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

• lazy_stack (bool, 可選) –

  引數應該密集堆疊還是惰性堆疊。預設為 False（密集堆疊）。

  注意

  lazy_stack 和 as_module 是互斥的特性。

  警告

  惰性和非惰性輸出之間的一個關鍵區別在於，非惰性輸出將使用所需的批次大小重新例項化引數，而 lazy_stack 只會將引數表示為惰性堆疊。這意味著，當 lazy_stack=True 時，可以將原始引數安全地傳遞給最佳化器，而設定為 True 時，則需要傳遞新引數。

  警告

  雖然使用惰性堆疊來保持原始引數引用可能很誘人，但請記住，每次呼叫 get() 時，惰性堆疊都會執行堆疊。這將需要記憶體（引數大小的 N 倍，如果構建圖形則更多）和時間來計算。這也意味著最佳化器將包含更多引數，並且像 step() 或 zero_grad() 這樣的操作將需要更長的時間才能執行。通常，lazy_stack 應僅保留給極少數用例。

• expand_identical (bool, 可選) – 如果為 True 並且正在將相同的引數（相同的標識）堆疊到自身，則將返回此引數的擴充套件版本。當 lazy_stack=True 時，將忽略此引數。

示例

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

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

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

classmethod from_namedtuple(named_tuple, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)

以遞迴方式將命名元組轉換為 TensorDict。

引數:

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

關鍵字引數:

• auto_batch_size (bool, 可選) – 如果 True，將自動計算批處理大小。 預設為 False。

• batch_dims (int, 可選) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每級完全批處理大小）。

• device (torch.device, 可選) – 將在其上建立 TensorDict 的裝置。預設為 None。

• batch_size (torch.Size, 可選) – 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_struct_array(struct_array: ndarray, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None) → T

將結構化的 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, 可選) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。預設為 None（每級完全批處理大小）。

• device (torch.device, optional) –

  將在其上建立 TensorDict 的裝置。預設為 None。

  注意

  更改裝置（即，指定除 None 或 "cpu" 之外的任何裝置）將傳輸資料，從而導致返回的資料的記憶體位置發生變化。

• batch_size (torch.Size, optional) – 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, 可選) – 如果 auto_batch_size 為 True，則定義輸出 tensordict 應具有的維度數。 預設為 None（每個級別的完整批處理大小）。

• device (torch.device, 可選) – 將在其上建立 TensorDict 的裝置。預設為 None。

• batch_size (torch.Size, 可選) – 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) → T

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

引數:

• dim (int) – 沿著該維度收集元素

• index (torch.Tensor) – 一個長張量，其維度數量與 tensordict 的維度數量匹配，只是兩個維度之間只有一個維度不同（即收集維度）。其元素引用沿所需維度收集的索引。

• out (TensorDictBase, optional) – 目標 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) → T | None

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

引數:

• dst (int) – 將呼叫 gather_and_stack() 的目標 worker 的 rank。

• group (torch.distributed.ProcessGroup, optional) – 如果設定，指定的程序組將被用於通訊。否則，預設程序組將被使用。預設為 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, *args, **kwargs) → Tensor

獲取與輸入鍵關聯的儲存值。

引數:

• 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)`。

示例

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

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

從鍵 key 中索引為 idx 的位置獲取 tensordict 的值。

引數:

• key (str, str 的元組) – 要檢索的鍵。

• index (int, slice, torch.Tensor, iterable) – 張量的索引。

• default (torch.Tensor) – 如果 tensordict 中不存在該鍵，則返回預設值。

返回值:

索引化的張量。

示例

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

get_item_shape(key: NestedKey)

返回條目的形狀，可能會避免呼叫 get()。

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

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

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

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

引數:

• key (NestedKey) – NonTensorData 物件的位置。

• default (Any, optional) – 如果找不到鍵，則返回的值。

返回：tensordict.tensorclass.NonTensorData 的內容，

或者，如果它不是 tensordict.tensorclass.NonTensorData，則返回與 key 對應的條目（如果找不到該條目，則返回 default）。

示例

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

property grad

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

half()

將所有張量轉換為 torch.half。

int()

將所有張量轉換為 torch.int。

int16()

將所有張量轉換為 torch.int16。

int32()

將所有張量轉換為 torch.int32。

int64()

將所有張量轉換為 torch.int64。

int8()

將所有張量轉換為 torch.int8。

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, optional) – 如果設定，指定的程序組將被用於通訊。否則，預設程序組將被使用。預設為 None。

• return_premature (bool) – 如果 True，則返回一個 future 列表，用於等待直到 tensordict 被更新。 預設為 False，即在該呼叫中等待直到更新完成。

• init_tag (int) – 源工作程序使用的 init_tag。

• pseudo_rand (bool) – 如果為 True，則標籤序列將是偽隨機的，允許從不同節點發送多個數據而不會重疊。 請注意，這些偽隨機數的生成成本很高（1e-5 秒/數字），這意味著它可能會降低演算法的執行時速度。 此值必須與傳遞給 isend() 的值匹配。 預設為 False。

返回值:

如果 return_premature=True，則返回一個 future 列表，用於等待

直到 tensordict 被更新。

is_consolidated()

檢查 TensorDict 是否具有合併的儲存。

is_contiguous()

返回一個布林值，指示所有張量是否是連續的。

is_empty() → bool

檢查 tensordict 是否包含任何葉節點。

is_memmap() → bool

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

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

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

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) → int

非同步傳送 tensordict 的內容。

引數:

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

關鍵字引數:

• group (torch.distributed.ProcessGroup, optional) – 如果設定，指定的程序組將被用於通訊。否則，預設程序組將被使用。預設為 None。

• init_tag (int) – 用於標記張量的初始標籤。 請注意，這將增加 TensorDict 中包含的張量數量。

• pseudo_rand (bool) – 如果為 True，則標籤序列將是偽隨機的，允許從不同節點發送多個數據而不會重疊。 請注意，這些偽隨機數的生成成本很高（1e-5 秒/數字），這意味著它可能會降低演算法的執行時速度。 預設為 False。

示例

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import multiprocessing as mp
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://: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() → T

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

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

isnan() → T

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

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

isneginf() → T

測試輸入的每個元素是否為負無窮大。

isposinf() → T

測試輸入的每個元素是否為負無窮大。

isreal() → T

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

items(include_nested: bool = False, leaves_only: bool = False, is_leaf=None, *, sort: bool = False) → Iterator[tuple[str, torch.Tensor]]

返回 tensordict 的鍵值對生成器。

引數:

• include_nested (bool, optional) – 如果為 True, 將返回巢狀值。 預設為 False。

• leaves_only (bool, optional) – 如果為 False, 將僅返回葉節點。 預設為 False。

• is_leaf (callable, optional) –

  一個可呼叫物件，作用於一個類型別，返回一個布林值，指示是否應將該類視為葉子節點。

  注意

  is_leaf 的目的不是阻止遞迴呼叫到巢狀的 tensordict 中，而是當 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(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) → _PersistentTDKeysView

返回 tensordict 鍵的生成器。

警告

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

引數:

• include_nested (bool, optional) – 如果為 True, 將返回巢狀值。 預設為 False。

• leaves_only (bool, optional) – 如果為 False, 將僅返回葉節點。 預設為 False。

• is_leaf (callable, optional) –

  一個可呼叫物件，作用於一個類型別，返回一個布林值，指示是否應將該類視為葉子節點。

  注意

  is_leaf 的目的不是阻止遞迴呼叫到巢狀的 tensordict 中，而是當 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=0, *, out=None, **kwargs)

建立一個 tensordicts 的惰性堆疊。

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

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

基於標量或張量 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, tensor or float) – 插值公式的權重。

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

lerp() 的原地版本。

lgamma() → T

計算 TensorDict 中每個元素的 lgamma() 值。

lgamma_() → T

原地 (in-place) 計算 TensorDict 中每個元素的 lgamma() 值。

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

從磁碟載入一個 tensordict。

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

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

從磁碟載入一個 tensordict 到當前的 tensordict 中。

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

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

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

引數:

• prefix (str 或 資料夾的 Path) – 儲存 tensordict 的資料夾路徑。

• device (torch.device 或 等價物, 可選) – 如果提供，資料將被非同步轉換為該裝置。支援 “meta” 裝置，在這種情況下，資料不會被載入，但會建立一組空的“meta”張量。這有助於瞭解總模型大小和結構，而無需實際開啟任何檔案。

• non_blocking (bool, 可選) – 如果 True，則在將張量載入到裝置後不會呼叫 synchronize。預設為 False。

• out (TensorDictBase, 可選) – 可選的 tensordict，資料將被寫入到該 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()

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

示例

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

tensordict 也可以載入到“meta”裝置上，或者作為偽張量 (fake tensor)。

示例

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

load_memmap_(prefix: str | pathlib.Path)

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

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

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

載入一個 state-dict (狀態字典)，格式如 state_dict()，到 tensordict 中。

引數:

• state_dict (OrderedDict) – 要複製的 state_dict。

• strict (bool, 可選) – 是否嚴格強制要求 state_dict 中的鍵與此 tensordict 的 torch.nn.Module.state_dict() 函式返回的鍵匹配。 預設值: True

• assign (bool, 可選) – 是否將狀態字典中的專案分配給 tensordict 中對應的鍵，而不是就地複製到 tensordict 的當前張量中。 當 False 時，保留當前模組中張量的屬性，而當 True 時，保留狀態字典中張量的屬性。 預設值: False

• from_flatten (bool, 可選) – 如果 True，則假定輸入 state_dict 已被展平。 預設為 False。

示例

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

lock_() → T

鎖定 tensordict，禁止非原地操作。

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

此方法可以用作裝飾器。

示例

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

log() → T

計算 TensorDict 中每個元素的 log() 值。

log10() → T

計算 TensorDict 中每個元素的 log10() 值。

log10_() → T

就地 (in-place) 計算 TensorDict 中每個元素的 log10() 值。

log1p() → T

計算 TensorDict 中每個元素的 log1p() 值。

log1p_() → T

就地 (in-place) 計算 TensorDict 中每個元素的 log1p() 值。

log2() → T

計算 TensorDict 中每個元素的 log2() 值。

log2_() → T

就地 (in-place) 計算 TensorDict 中每個元素的 log2() 值。

log_() → T

就地 (in-place) 計算 TensorDict 中每個元素的 log() 值。

logical_and(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) → TensorDictBase

在 self 和 other 之間執行邏輯與 (AND) 操作。

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

引數:

other (TensorDictBase 或 torch.Tensor) – 用於執行邏輯與操作的 tensor 或 TensorDict。

關鍵字引數:

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

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

返回給定維度 dim 中，輸入 tensordict 每行指數和的對數。 該計算經過數值穩定處理。

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

引數:

• dim (int 或 int 元組, 可選) – 要縮減的維度。如果為 None，則 tensordict 的所有批處理維度都將被縮減。

• keepdim (bool) – 輸出 tensordict 是否保留 dim。

關鍵字引數:

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

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

根據給定的形狀和可能的 dtype 建立一個空的記憶體對映 tensor。

警告

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

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

引數:

• key (NestedKey) – 要寫入的新條目的鍵。 如果該鍵已存在於 tensordict 中，則會引發異常。

• shape (torch.Size 或等效值, 巢狀 tensor 的情況是 torch.Tensor) – 要寫入的 tensor 的形狀。

關鍵字引數:

dtype (torch.dtype, 可選) – 新 tensor 的 dtype。

返回值:

一個新的記憶體對映 tensor。

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

根據給定的儲存，形狀和可能的 dtype 建立一個空的記憶體對映 tensor。

警告

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

注意

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

引數:

• key (NestedKey) – 要寫入的新條目的鍵。 如果該鍵已存在於 tensordict 中，則會引發異常。

• storage (torch.UntypedStorage) – 用於新 MemoryMappedTensor 的儲存。 必須是物理記憶體儲存。

• shape (torch.Size 或等效值, 巢狀 tensor 的情況是 torch.Tensor) – 要寫入的 tensor 的形狀。

關鍵字引數:

dtype (torch.dtype, 可選) – 新 tensor 的 dtype。

返回值:

具有給定儲存的新記憶體對映 tensor。

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, optional) – 如果 False，則新張量將共享輸入的元資料（如形狀和 dtype），但內容將為空。預設為 True。

返回值:

具有給定儲存的新記憶體對映 tensor。

map(fn: Callable, dim: int = 0, num_workers: Optional[int] = None, *, out: Optional[TensorDictBase] = 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, index_with_generator: bool = False, pbar: bool = False, mp_start_method: Optional[str] = None)

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

此方法將透過將 tensordict 分割成大小相等的 tensordict，並將操作分派到所需數量的 worker 上來將函式應用於 tensordict 例項。

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

注意

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

引數:

• fn (callable) – 要應用於 tensordict 的函式。支援類似於 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的簽名。

• dim (int, optional) – tensordict 將沿其分塊的維度。

• num_workers (int, optional) – worker 的數量。與 pool 互斥。如果未提供，worker 的數量將設定為可用的 CPU 數量。

關鍵字引數:

• out (TensorDictBase, optional) – 輸出的可選容器。它沿提供的 dim 的批次大小必須與 self.ndim 匹配。如果它是共享的或記憶體對映的（is_shared() 或 is_memmap() 返回 True），它將在遠端程序中填充，避免資料向內傳輸。否則，self 切片的資料將被髮送到程序，在當前程序上收集並就地寫入 out。

• chunksize (int, 可選) – 每個資料塊的大小。如果 chunksize 為 0，則會沿所需維度解綁 tensordict，並在應用函式後重新堆疊它。而如果 chunksize>0，則會分割 tensordict 並呼叫 torch.cat() 函式來連線結果 tensordict 列表。 如果未提供，則塊數將等於 worker 的數量。 對於非常大的 tensordict，如此大的塊可能無法放入記憶體中進行操作，並且可能需要更多塊才能使操作切實可行。 此引數與 num_chunks 互斥。

• num_chunks (int, 可選) – 將 tensordict 分割成的塊數。 如果未提供，則塊數將等於 worker 的數量。 對於非常大的 tensordict，如此大的塊可能無法放入記憶體中進行操作，並且可能需要更多塊才能使操作切實可行。 此引數與 chunksize 互斥。

• pool (mp.Pool, 可選) – 用於執行作業的多程序 Pool 例項。如果未提供，將在 map 方法中建立一個池。

• generator (torch.Generator, 可選) –

  用於播種的生成器。 將從中生成一個基本種子，並且 pool 的每個 worker 將以提供的種子遞增一個從 0 到 num_workers 的唯一整數來播種。 如果未提供生成器，則將使用隨機整數作為種子。 要與未播種的 worker 一起使用，應單獨建立一個池並直接傳遞給 map()。

  注意

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

  注意

  worker 播種的目的是使每個 worker 都有獨立的種子，而不是使 map 方法的呼叫中產生可重複的結果。換句話說，兩個實驗可能會並且很可能會返回不同的結果，因為不可能知道哪個 worker 將選擇哪個作業。但是，我們可以確保每個 worker 都有不同的種子，並且每個 worker 上的偽隨機操作都將不相關。

• max_tasks_per_child (int, 可選) – 每個子程序選擇的最大作業數。預設為 None，即對作業數量沒有限制。

• worker_threads (int, 可選) – worker 的執行緒數。預設為 1。

• index_with_generator (bool, 可選) – 如果為 True，則將在查詢期間完成 tensordict 的拆分/分塊，從而節省初始化時間。 請注意，chunk() 和 split() 比索引（在生成器中使用）更有效，因此在初始化時獲得的處理時間可能會對總執行時間產生負面影響。 預設為 False。

• pbar (bool, 可選) – 如果為 True，將顯示進度條。需要安裝 tqdm。預設為 False。

• mp_start_method (str, 可選) – 用於多程序的啟動方法。 如果未提供，將使用預設的啟動方法。 接受的字串為 "fork" 和 "spawn"。 請記住，"cuda" 張量不能在使用 "fork" 啟動方法的程序之間共享。 如果將 pool 傳遞給 map 方法，則這不起作用。

示例

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

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

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

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

此方法將透過將 tensordict 例項分割成大小相等的 tensordict 塊，並將操作分派到所需數量的工作程序上，從而將函式應用於 tensordict 例項。它將一次生成一個結果。

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

注意

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

注意

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

引數:

• fn (callable) – 要應用於 tensordict 的函式。支援類似於 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的簽名。

• dim (int, optional) – tensordict 將沿其分塊的維度。

• num_workers (int, optional) – worker 的數量。與 pool 互斥。如果未提供，worker 的數量將設定為可用的 CPU 數量。

關鍵字引數:

• shuffle (bool, 可選) – 指示是否應全域性打亂索引。如果 True，則每個批次將包含不連續的樣本。 如果 index_with_generator=False 並且 shuffle=True`，則會引發錯誤。 預設為 False。

• chunksize (int, 可選) – 每個資料塊的大小。如果 chunksize 為 0，則會沿所需維度解綁 tensordict，並在應用函式後重新堆疊它。而如果 chunksize>0，則會分割 tensordict 並呼叫 torch.cat() 函式來連線結果 tensordict 列表。 如果未提供，則塊數將等於 worker 的數量。 對於非常大的 tensordict，如此大的塊可能無法放入記憶體中進行操作，並且可能需要更多塊才能使操作切實可行。 此引數與 num_chunks 互斥。

• num_chunks (int, 可選) – 將 tensordict 分割成的塊數。 如果未提供，則塊數將等於 worker 的數量。 對於非常大的 tensordict，如此大的塊可能無法放入記憶體中進行操作，並且可能需要更多塊才能使操作切實可行。 此引數與 chunksize 互斥。

• pool (mp.Pool, 可選) – 用於執行作業的多程序 Pool 例項。如果未提供，將在 map 方法中建立一個池。

• generator (torch.Generator, 可選) –

  用於播種的生成器。 將從中生成一個基本種子，並且 pool 的每個 worker 將以提供的種子遞增一個從 0 到 num_workers 的唯一整數來播種。 如果未提供生成器，則將使用隨機整數作為種子。 要與未播種的 worker 一起使用，應單獨建立一個池並直接傳遞給 map()。

  注意

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

  注意

  worker 播種的目的是使每個 worker 都有獨立的種子，而不是使 map 方法的呼叫中產生可重複的結果。換句話說，兩個實驗可能會並且很可能會返回不同的結果，因為不可能知道哪個 worker 將選擇哪個作業。但是，我們可以確保每個 worker 都有不同的種子，並且每個 worker 上的偽隨機操作都將不相關。

• max_tasks_per_child (int, 可選) – 每個子程序選擇的最大作業數。預設為 None，即對作業數量沒有限制。

• worker_threads (int, 可選) – worker 的執行緒數。預設為 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, 可選) – 用於多程序的啟動方法。 如果未提供，將使用預設的啟動方法。 接受的字串為 "fork" 和 "spawn"。 請記住，"cuda" 張量不能在使用 "fork" 啟動方法的程序之間共享。 如果將 pool 傳遞給 map 方法，則這不起作用。

示例

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

masked_fill(mask, value)

masked_fill 的非原地版本。

引數:

• mask (布林型 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_(mask, value)

使用所需的值填充與掩碼對應的值。

引數:

• mask (布林型 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) → T

遮蔽 TensorDict 的所有張量，並返回一個新的 TensorDict 例項，該例項具有指向遮蔽值的類似鍵。

引數:

mask (torch.Tensor) – 用於張量的布林掩碼。 形狀必須與 TensorDict batch_size 匹配。

示例

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

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

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

引數:

• dim (int, optional) – 如果 None, 返回一個包含所有葉子節點的最大值的無維度 tensordict (如果可以計算). 如果是整數，max 將僅在指定維度與 tensordict 形狀相容時，在該維度上呼叫。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上進行縮減，並返回單個縮減張量。預設為 False。

• return_argmins (bool, optional) – 當傳遞 dim 引數時，max() 返回一個包含 values 和 indices 的命名元組。與此等效的 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.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) → T

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

引數:

other (TensorDict or Tensor) – 另一個輸入 tensordict 或 tensor。

關鍵字引數:

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

maximum_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

maximum() 的原地版本。

注意

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

classmethod maybe_dense_stack(input, dim=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) → tensordict.base.TensorDictBase | torch.Tensor

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

引數:

• dim (int, int 的元組, str, optional) – 如果 None，返回一個包含所有葉子節點的平均值的無維度 tensordict (如果可以計算)。如果是整數或整數元組，mean 將僅在指定維度與 tensordict 形狀相容時，在該維度上呼叫。目前只允許使用字串 “feature”。使用 dim=”feature” 將實現對所有特徵維度的縮減。如果 reduce=True，則將返回一個與 TensorDict 的批次大小形狀相同的張量。否則，將返回一個具有與 self 相同結構的新 tensordict，但具有縮減的特徵維度。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• dtype (torch.dtype, optional) – 返回張量的所需資料型別。如果指定，則在執行操作之前將輸入張量強制轉換為 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) → T

將所有張量寫入到新 tensordict 中對應的記憶體對映張量。

引數:

• prefix (str) – 儲存記憶體對映張量的目錄字首。目錄樹結構將模仿 tensordict 的結構。

• copy_existing (bool) – 如果為 False（預設值），則當 tensordict 中的條目已經是儲存在磁碟上且有關聯檔案的張量，但未根據字首儲存在正確位置時，將引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

• num_threads (int, optional) – 用於寫入 memmap 張量的執行緒數。預設為 0。

• return_early (bool, optional) – 如果 True 並且 num_threads>0，則該方法將返回 tensordict 的 future。

• share_non_tensor (bool, optional) – 如果 True，則非張量資料將在程序和寫入操作（例如，就地更新或設定）之間共享，單個節點內任何 worker 上的寫入操作都將更新所有其他 worker 上的值。如果非張量葉的數量很多（例如，共享大量非張量資料），則可能導致 OOM 或類似錯誤。預設為 False。

• existsok (bool, optional) – 如果 False，則如果張量已存在於同一路徑中，將引發異常。預設為 True。

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

返回值:

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

注意

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

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

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

引數:

• prefix (str) – 儲存記憶體對映張量的目錄字首。目錄樹結構將模仿 tensordict 的結構。

• copy_existing (bool) – 如果為 False（預設值），則當 tensordict 中的條目已經是儲存在磁碟上且有關聯檔案的張量，但未根據字首儲存在正確位置時，將引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

• num_threads (int, optional) – 用於寫入 memmap 張量的執行緒數。預設為 0。

• return_early (bool, optional) – 如果為 True 且 num_threads>0，則該方法將返回 tensordict 的 future。可以使用 future.result() 查詢結果 tensordict。

• share_non_tensor (bool, optional) – 如果 True，則非張量資料將在程序和寫入操作（例如，就地更新或設定）之間共享，單個節點內任何 worker 上的寫入操作都將更新所有其他 worker 上的值。如果非張量葉的數量很多（例如，共享大量非張量資料），則可能導致 OOM 或類似錯誤。預設為 False。

• existsok (bool, optional) – 如果 False，則如果張量已存在於同一路徑中，將引發異常。預設為 True。

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

返回值:

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

注意

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

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

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

引數:

• prefix (str) – 儲存記憶體對映張量的目錄字首。目錄樹結構將模仿 tensordict 的結構。

• copy_existing (bool) – 如果為 False（預設值），則當 tensordict 中的條目已經是儲存在磁碟上且有關聯檔案的張量，但未根據字首儲存在正確位置時，將引發異常。如果為 True，則任何現有張量都將被複制到新位置。

關鍵字引數:

• num_threads (int, optional) – 用於寫入 memmap 張量的執行緒數。預設為 0。

• return_early (bool, optional) – 如果 True 並且 num_threads>0，則該方法將返回 tensordict 的 future。

• share_non_tensor (bool, optional) – 如果 True，則非張量資料將在程序和寫入操作（例如，就地更新或設定）之間共享，單個節點內任何 worker 上的寫入操作都將更新所有其他 worker 上的值。如果非張量葉的數量很多（例如，共享大量非張量資料），則可能導致 OOM 或類似錯誤。預設為 False。

• existsok (bool, optional) – 如果 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) → TensorDictBase | torch.Tensor

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

引數:

• dim (int, 可選) – 如果 None, 返回一個包含所有葉子節點的最小值（如果可以計算）的無維度 tensordict。 如果是整數，則僅當此維度與 tensordict 形狀相容時，才在指定的維度上呼叫 min。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• reduce (bool, optional) – 如果為 True，則將在所有 TensorDict 值上進行縮減，並返回單個縮減張量。預設為 False。

• return_argmins (bool, 可選) – min() 在傳遞 dim 引數時返回一個包含值和索引的命名元組。 等效的 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.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.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) → T

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

引數:

other (TensorDict or Tensor) – 另一個輸入 tensordict 或 tensor。

關鍵字引數:

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

minimum_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

minimum() 的原地版本。

注意

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

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

將 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.base.TensorDictBase | torch.Tensor) → T

mul() 的原地版本。

注意

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

named_apply(fn: Callable, *others: T, nested_keys: bool = False, 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: Optional[bool] = None, propagate_lock: bool = False, call_on_nested: bool = False, out: Optional[TensorDictBase] = None, **constructor_kwargs) → Optional[T]

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

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

引數:

• fn (Callable) – 要應用於 tensordict 中的 (name, tensor) 對的函式。 對於每個葉節點，只會使用其葉節點名稱（而不是完整的 NestedKey）。

• *others (TensorDictBase 例項, 可選) – 如果提供，這些 tensordict 例項應具有與 self 匹配的結構。 fn 引數應接收與 tensordict 數量一樣多的未命名輸入，包括 self。如果其他 tensordict 缺少條目，則可以透過 default 關鍵字引數傳遞預設值。

• nested_keys (bool, optional) – 如果 True，將使用到葉節點的完整路徑。 預設為 False，即只有最後一個字串會傳遞給函式。

• batch_size (int 序列, 可選) – 如果提供，生成的 TensorDict 將具有所需的 batch_size。 batch_size 引數應與轉換後的 batch_size 匹配。這是一個僅關鍵字引數。

• device (torch.device, 可選) – 生成的裝置，如果存在。

• names (str 列表, 可選) – 新的維度名稱，在 batch_size 被修改的情況下。

• inplace (bool, optional) – 如果為 True，則進行原地更改。預設為 False。這是一個僅限關鍵字的引數。

• default (Any, 可選) – 其他 tensordict 中缺少條目的預設值。 如果未提供，缺少條目將引發 KeyError。

• filter_empty (bool, optional) – 如果 True，將過濾掉空的 tensordict。這也會降低計算成本，因為不會建立和銷燬空資料結構。為了向後相容，預設為 False。

• propagate_lock (bool, optional) – 如果 True，鎖定的 tensordict 將生成另一個鎖定的 tensordict。預設為 False。

• call_on_nested (bool, optional) –

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

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

• out (TensorDictBase, 可選) –

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

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

  警告

  如果對 tensordict 執行的操作需要訪問多個鍵才能進行單次計算，則提供等於 self 的 out 引數可能會導致該操作默默地產生錯誤的結果。 例如

  >>> td = TensorDict({"a": 1, "b": 1})
  >>> td.apply(lambda x: x+td["a"])["b"] # Right!
  tensor(2)
  >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong!
  tensor(3)
  

• **constructor_kwargs – 要傳遞給 TensorDict 建構函式的其他關鍵字引數。

返回值:

具有 transformed_in 張量的新 tensordict。

示例

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

注意

如果函式返回 None，則忽略該條目。 這可用於過濾 tensordict 中的資料

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

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) → tensordict.base.TensorDictBase | torch.Tensor

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

引數:

• dim (int, tuple of int, optional) – 如果 None，返回一個無量綱的 tensordict，其中包含所有葉子的平均值（如果可以計算）。 如果是整數或整數元組，則僅當此維度與 tensordict 形狀相容時，才會在指定的維度上呼叫 mean。 目前只允許使用 “feature” 字串。 使用 dim=”feature” 將實現對所有特徵維度的縮減。 如果 reduce=True，則將返回一個形狀為 TensorDict 批大小的張量。 否則，將返回一個新的 tensordict，其結構與 self 相同，並且減少了特徵維度。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• dtype (torch.dtype, optional) – 返回張量的所需資料型別。如果指定，則在執行操作之前將輸入張量強制轉換為 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) → tensordict.base.TensorDictBase | torch.Tensor

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

引數:

• dim (int, int 元組, 可選) – 如果為 None，則返回一個無維度 tensordict，其中包含所有葉節點的總和值（如果可以計算）。 如果是整數或整數元組，則僅當此維度與 tensordict 形狀相容時，才對指定維度呼叫 sum。 目前只允許使用 “feature” 字串。 使用 dim=”feature” 將實現對所有特徵維度的縮減。 如果 reduce=True，將返回一個形狀與 TensorDict 的批次大小相同的張量。 否則，將返回一個新的 tensordict，其結構與 self 相同，但特徵維度已縮減。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• dtype (torch.dtype, optional) – 返回張量的所需資料型別。如果指定，則在執行操作之前將輸入張量強制轉換為 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() → T

計算 TensorDict 中每個元素的 neg() 值。

neg_() → T

就地計算 TensorDict 中每個元素的 neg() 值。

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

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

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

引數:

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

關鍵字引數:

• dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。 預設值：如果為 None，則 torch.dtype 將保持不變。

• device (torch.device, 可選) – 返回的 tensordict 的所需裝置。 預設值：如果為 None，則 torch.device 將保持不變。

• requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。 預設值：False。

• layout (torch.layout, 可選) – 返回的 TensorDict 值的所需佈局。 預設值：torch.strided。

• pin_memory (bool, 可選) – 如果設定，返回的張量將在鎖頁記憶體中分配。 僅適用於 CPU 張量。 預設值：False。

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

返回一個大小為 size 且填充了指定值的 TensorDict。

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

引數:

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

• fill_value (標量) – 用於填充輸出張量的數值。

關鍵字引數:

• dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。 預設值：如果為 None，則 torch.dtype 將保持不變。

• device (torch.device, 可選) – 返回的 tensordict 的所需裝置。 預設值：如果為 None，則 torch.device 將保持不變。

• requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。 預設值：False。

• layout (torch.layout, 可選) – 返回的 TensorDict 值的所需佈局。 預設值：torch.strided。

• pin_memory (bool, 可選) – 如果設定，返回的張量將在鎖頁記憶體中分配。 僅適用於 CPU 張量。 預設值：False。

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

返回一個大小為 size 且填充了指定值的 TensorDict。

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

引數:

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

關鍵字引數:

• dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。 預設值：如果為 None，則 torch.dtype 將保持不變。

• device (torch.device, 可選) – 返回的 tensordict 的所需裝置。 預設值：如果為 None，則 torch.device 將保持不變。

• requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。 預設值：False。

• layout (torch.layout, 可選) – 返回的 TensorDict 值的所需佈局。 預設值：torch.strided。

• pin_memory (bool, 可選) – 如果設定，返回的張量將在鎖頁記憶體中分配。 僅適用於 CPU 張量。 預設值：False。

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

返回一個以張量 data 作為資料的新 TensorDict。

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

data 也可以是張量集合（TensorDict 或 tensorclass），在這種情況下，new_tensor 方法會迭代 self 和 data 的張量對。

引數:

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

關鍵字引數:

• dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。 預設值：如果為 None，則 torch.dtype 將保持不變。

• device (torch.device, 可選) – 返回的 tensordict 的所需裝置。 預設值：如果為 None，則 torch.device 將保持不變。

• requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。 預設值：False。

• pin_memory (bool, 可選) – 如果設定，返回的張量將在鎖頁記憶體中分配。 僅適用於 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)

返回一個大小為 size 的 TensorDict，並用 0 填充。

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

引數:

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

關鍵字引數:

• dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。 預設值：如果為 None，則 torch.dtype 將保持不變。

• device (torch.device, 可選) – 返回的 tensordict 的所需裝置。 預設值：如果為 None，則 torch.device 將保持不變。

• requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。 預設值：False。

• layout (torch.layout, 可選) – 返回的 TensorDict 值的所需佈局。 預設值：torch.strided。

• pin_memory (bool, 可選) – 如果設定，返回的張量將在鎖頁記憶體中分配。 僅適用於 CPU 張量。 預設值：False。

non_tensor_items(include_nested: bool = False)

返回所有非 tensor 的葉節點，可以遞迴。

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

計算 tensordict 中每個 tensor 的範數。

關鍵字引數:

• out (TensorDict, optional) – 輸出 tensordict。

• dtype (torch.dtype, optional) – 輸出 dtype (torch>=2.4)。

numel() → int

批次中的元素總數。

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

numpy()

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

非 tensor 資料按原樣公開。

示例

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

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

計算引數的數量（可索引專案的總數），僅考慮 tensor。

關鍵字引數:

count_duplicates (bool) – 是否將重複的張量視為獨立的。 如果 False，則只會丟棄嚴格相同的張量（來自公共基本張量的相同檢視但不同的 ID 將被計算兩次）。 預設為 True（假設每個張量都是單個副本）。

permute(*args, **kwargs)

返回 tensordict 的檢視，其中批次維度根據 dims 排列。

引數:

• *dims_list (int) – tensordict 批次維度的新順序。或者，可以提供單個整數迭代器。

• dims (list of int) – 呼叫 permute(…) 的另一種方式。

返回值:

一個新的 tensordict，批次維度按所需的順序排列。

示例

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

pin_memory(*args, **kwargs)

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

引數:

• num_threads (int or str) – 如果提供，用於呼叫葉子節點上的 pin_memory 的執行緒數。 預設為 None，這將在 ThreadPoolExecutor(max_workers=None) 中設定大量的執行緒。 要在主執行緒上執行對 pin_memory() 的所有呼叫，請傳遞 num_threads=0。

• inplace (bool, optional) – 如果 True，則就地修改 tensordict。 預設為 False。

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

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

引數:

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

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

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

如果該值不存在且沒有提供預設值，則會丟擲一個 KeyError。

引數:

• key (str 或 巢狀鍵) – 要查詢的條目。

• default (Any, 可選) – 如果找不到鍵，則返回的值。

示例

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

popitem() → Tuple[NestedKey, Tensor]

移除最後插入到 TensorDict 中的項。

popitem 只會返回非巢狀的值。

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

計算 self 中每個元素的 other 次冪，並返回包含結果的張量。

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

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

引數:

other (float, tensor 或 tensordict) – 指數值

關鍵字引數:

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

pow_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

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) → tensordict.base.TensorDictBase | torch.Tensor

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

引數:

• dim (int, int 元組, 可選) – 如果 None，則返回包含所有葉子的乘積值的無維度 tensordict（如果可以計算）。如果是整數或整數元組，則僅當此維度與 tensordict 形狀相容時，才在指定維度上呼叫 prod。目前僅允許使用 “feature” 字串。 使用 dim="feature" 將實現對所有特徵維度的縮減。 如果 reduce=True，則將返回 TensorDict 批大小形狀的張量。 否則，將返回與 self 具有相同結構但減少了特徵維度的新 tensordict。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• dtype (torch.dtype, optional) – 返回張量的所需資料型別。如果指定，則在執行操作之前將輸入張量強制轉換為 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()

將所有張量轉換為 torch.qint32。

qint8()

將所有張量轉換為 torch.qint8。

quint4x2()

將所有張量轉換為 torch.quint4x2。

quint8()

將所有張量轉換為 torch.quint8。

reciprocal() → T

計算 TensorDict 中每個元素的 reciprocal() 值。

reciprocal_() → T

計算 TensorDict 中每個元素的 reciprocal() 值，並直接修改 TensorDict。

record_stream(stream: Stream) → T

將 tensordict 標記為已被該流使用。

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

有關更多資訊，請參閱 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, optional) – 如果設定，指定的程序組將被用於通訊。否則，預設程序組將被使用。預設為 None。

• init_tag (int) – 源工作程序使用的 init_tag。

• pseudo_rand (bool) – 如果為 True，則標籤序列將是偽隨機的，允許從不同節點發送多個數據而不會重疊。 請注意，生成這些偽隨機數成本很高 (1e-5 秒/數字)，這意味著它可能會減慢演算法的執行時間。 此值必須與傳遞給 send() 的值匹配。 預設為 False。

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

跨所有機器縮減 tensordict。

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

refine_names(*names) → T

根據名稱最佳化 self 的維度名稱。

最佳化是一種重新命名的特殊情況，可以“提升”未命名的維度。 可以將 None dim 最佳化為具有任何名稱； 命名的 dim 只能最佳化為具有相同的名稱。

由於命名的張量可以與未命名的張量共存，因此最佳化名稱提供了一種很好的方式來編寫具有命名張量意識的程式碼，該程式碼可以與命名的和未命名的張量一起使用。

names 最多可以包含一個省略號 (…)。 省略號會貪婪地展開； 它會就地展開，以使用 self.names 中相應索引的名稱將 names 填充到與 self.dim() 相同的長度。

返回值：具有根據輸入命名的維度的相同 tensordict。

示例

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

rename(*names, **rename_map)

返回已重新命名維度的 tensordict 的克隆。

示例

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

rename_(*names, **rename_map)

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

示例

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

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

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

引數:

• old_key (str 或 巢狀鍵) – 要重新命名的鍵。

• new_key (str 或 巢狀鍵) – 條目的新名稱。

• safe (bool, 可選) – 如果 True，則當新鍵已存在於 TensorDict 中時，會引發錯誤。

返回值:

self

repeat(*repeats: int) → TensorDictBase

沿指定維度重複此張量。

與 expand() 不同，此函式會複製張量的資料。

警告

repeat() 的行為與 repeat() 不同，但更類似於 numpy.tile()。 對於類似於 numpy.repeat() 的運算子，請參閱 repeat_interleave()。

引數:

repeat (torch.Size, int..., int 元組 或 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(repeats: torch.Tensor | int, dim: Optional[int] = None, *, output_size: Optional[int] = None) → T

重複 TensorDict 的元素。

警告

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

引數:

• repeats (torch.Tensor 或 int) – 每個元素的重複次數。repeats 將被廣播以適應給定軸的形狀。

• dim (int, 可選) – 沿其重複值的維度。預設情況下，使用扁平化的輸入陣列，並返回一個扁平化的輸出陣列。

關鍵字引數:

output_size (int, 可選) – 給定軸的總輸出大小（例如，重複次數的總和）。如果給定，它將避免計算 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) → T

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

返回此 tensordict。

引數:

requires_grad (bool, 可選) – autograd 是否應該記錄此 tensordict 上的操作。 預設為 True。

reshape(*args, **kwargs) → T

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

引數:

*shape (int) – 結果 tensordict 的新形狀。

返回值:

具有重塑鍵的 TensorDict

示例

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

round() → T

計算 TensorDict 中每個元素的 round() 值。

round_() → T

就地計算 TensorDict 中每個元素的 round() 值。

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

將 tensordict 儲存到磁碟。

此函式是 memmap() 的代理。

property saved_path

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

只要 is_memmap() 返回 False（例如，當 tensordict 被解鎖時），此引數就會消失。

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

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

這些值不會被複制：對原始或新 tensordict 的張量進行就地修改會導致兩個 tensordict 都發生更改。

引數:

• *keys (str) – 要選擇的鍵

• inplace (bool) – 如果為 True，則就地修剪 tensordict。 預設值為 False。

• strict (bool, 可選) – 選擇不存在的鍵是否會返回錯誤。 預設值：True。

返回值:

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

注意

要選擇 tensordict 中的鍵並返回一個刪除了這些鍵的版本，請參閱 split_keys() 方法。

示例

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

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

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

引數:

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

關鍵字引數:

• group (torch.distributed.ProcessGroup, optional) – 如果設定，指定的程序組將被用於通訊。否則，預設程序組將被使用。預設為 None。

• init_tag (int) – 用於標記張量的初始標籤。 請注意，這將增加 TensorDict 中包含的張量數量。

• pseudo_rand (bool) – 如果為 True，則標籤序列將是偽隨機的，允許從不同節點發送多個數據而不會重疊。 請注意，這些偽隨機數的生成成本很高（1e-5 秒/數字），這意味著它可能會降低演算法的執行時速度。 預設為 False。

示例

>>> from torch import multiprocessing as mp
>>> from tensordict import TensorDict
>>> import torch
>>>
>>>
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://: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) → T

將指定的鍵從 tensordict 中就地分離。

參見

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

參見

此方法等效於呼叫 exclude()，不同之處在於它返回資料的另一個拆分。

引數:

• keys (NestedKey) – 要從 tensordict 中分離的鍵。

• default (Any, optional) – 缺少鍵時要返回的值。 如果未指定且 strict=True，則會引發異常。 否則，任何缺失鍵的預設值將為 None，除非另有說明。

• strict (bool, optional) – 如果 True，則在缺少鍵時會引發異常。 預設為 True。

• filter_empty (bool, optional) – 如果 True，則將刪除 self 中的空 tensordict。 預設為 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: Tensor, inplace: bool = False, *, non_blocking: bool = False, **kwargs: Any) → T

設定新的鍵值對。

引數:

• key (str, str 元組) – 要設定的鍵的名稱。

• item (torch.Tensor 或 等效項, TensorDictBase 例項) – 要儲存在 tensordict 中的值。

• inplace (bool, optional) – 如果 True 且鍵與 tensordict 中的現有鍵匹配，則該鍵值對將就地更新。 如果 inplace 為 True 且找不到條目，則會新增它。 對於更嚴格的就地操作，請改用 set_()。 預設為 False。

關鍵字引數:

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

返回值:

self

示例

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

set_(key: NestedKey, item: Tensor, *, non_blocking: bool = False) → T

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

引數:

• key (str) – 值的名稱

• item (torch.Tensor 或 相容型別, TensorDictBase) – 要儲存在 tensordict 中的值

關鍵字引數:

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

返回值:

self

示例

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

set_at_(key: NestedKey, value: Tensor, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], *, non_blocking: bool = False) → T

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

引數:

• key (str, str 元組) – 要修改的鍵。

• value (torch.Tensor) – 要在索引 index 處設定的值

• index (int, tensor 或 tuple) – 用於寫入值的索引。

關鍵字引數:

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

返回值:

self

示例

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

set_non_tensor(key: NestedKey, value: Any)

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

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

return: self

示例

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

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

如果 key 不在 tensordict 中，則使用 default 的值插入 key 條目。

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

引數:

• key (str 或 nested key) – 值的名稱。

• default (torch.Tensor 或 相容型別, TensorDictBase) – 如果鍵尚未存在，則儲存在 tensordict 中的值。

返回值:

tensordict 中鍵的值。 如果先前未設定該鍵，則將為 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_()

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

然後 TensorDict 被鎖定，這意味著任何非就地寫入操作都將引發異常（例如，重新命名、設定或刪除條目）。 相反，一旦 tensordict 被解鎖，share_memory 屬性將變為 False，因為不再保證跨程序標識。

返回值:

self

sigmoid() → T

計算 TensorDict 的每個元素的 sigmoid() 值。

sigmoid_() → T

就地計算 TensorDict 的每個元素的 sigmoid() 值。

sign() → T

計算 TensorDict 的每個元素的 sign() 值。

sign_() → T

就地計算 TensorDict 的每個元素的 sign() 值。

sin() → T

計算 TensorDict 的每個元素的 sin() 值。

sin_() → T

就地計算 TensorDict 的每個元素的 sin() 值。

sinh() → T

計算 TensorDict 的每個元素的 sinh() 值。

sinh_() → T

就地計算 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 或 整數元組) – 將在計算 softmax 的 tensordict 維度。

• dtype (torch.dtype, 可選) – 返回 tensor 的所需資料型別。 如果指定，則在執行操作之前將輸入 tensor 強制轉換為 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 或 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, 可選) – 如果為 True，則從 self 中原地刪除鍵。預設為 False。

• default (Any, 可選) – 當鍵缺失時返回的值。如果未指定且 strict=True，則會引發異常。

• strict (bool, optional) – 如果 True，則在缺少鍵時會引發異常。 預設為 True。

• reproduce_struct (bool, 可選) – 如果為 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()

計算 self 的逐元素平方根。

sqrt_()

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=0, *, out=None)

沿著給定的維度將 tensordict 堆疊到一個 tensordict 中。

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

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

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

引數:

dim (int, 可選) – 條目應沿其堆疊的維度。

關鍵字引數:

• sorted (bool 或 NestedKeys 列表) – 如果為 True，則條目將按字母順序堆疊。如果為 False (預設)，則將使用字典順序。或者，可以提供鍵名列表，並且將相應地堆疊張量。這會產生一些開銷，因為將根據 tensordict 中的葉名稱列表檢查鍵列表。

• out (torch.Tensor, 可選) – 用於堆疊操作的可選目標張量。

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

將條目堆疊到一個新條目中，並可能刪除原始值。

引數:

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

關鍵詞引數

out_key (NestedKey): 堆疊輸入的新鍵名。 keep_entries (bool, 可選): 如果 False，則 keys 中的條目將被刪除。

預設為 False。

dim (int, 可選): 堆疊必須發生的維度。

預設為 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(destination=None, prefix='', keep_vars=False, flatten=False) → OrderedDict[str, Any]

從 tensordict 生成一個 state_dict。

除非將 flatten 設定為 True，否則 state-dict 的結構仍然是巢狀的。

一個 tensordict state-dict 包含重建 tensordict 所需的所有張量和元資料（當前不支援名稱）。

引數:

• destination (dict, 可選) – 如果提供，tensordict 的狀態將被更新到 dict 中，並返回相同的物件。否則，將建立並返回一個 OrderedDict。預設值: None。

• prefix (str, 可選) – 新增到張量名稱的字首，用於組合 state_dict 中的鍵。預設值: ''。

• keep_vars (bool, 可選) – 預設情況下，state dict 中返回的 torch.Tensor 項與 autograd 分離。 如果設定為 True，則不會執行分離。 預設值: False。

• flatten (bool, 可選) – 結構是否應使用 "." 字元展平。 預設為 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) → tensordict.base.TensorDictBase | torch.Tensor

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

引數:

• dim (int, int 元組, 可選) – 如果 None, 則返回一個無量綱的 tensordict，其中包含所有葉子的總和值（如果可以計算）。如果為整數或整數元組，則僅當此維度與 tensordict 形狀相容時，才對指定的維度呼叫 std。目前只允許使用 “feature” 字串。使用 dim="feature" 將實現所有特徵維度的縮減。如果 reduce=True，將返回一個形狀為 TensorDict 批大小的張量。否則，將返回一個新的 tensordict，其結構與 self 相同，但特徵維度已縮減。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• correction (int) – 樣本大小和樣本自由度之間的差異。預設為貝塞爾校正，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: str | torch.Tensor | None = None)

從 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 將用於操作雙方的所有缺失條目。

sub_(other: tensordict.base.TensorDictBase | torch.Tensor | float, alpha: Optional[float] = None)

sub() 的原地版本。

注意

In-place sub 不支援 default 關鍵字引數。

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

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

引數:

• dim (int, int 元組, 可選) – 如果為 None，則返回一個無維度 tensordict，其中包含所有葉節點的總和值（如果可以計算）。 如果是整數或整數元組，則僅當此維度與 tensordict 形狀相容時，才對指定維度呼叫 sum。 目前只允許使用 “feature” 字串。 使用 dim=”feature” 將實現對所有特徵維度的縮減。 如果 reduce=True，將返回一個形狀與 TensorDict 的批次大小相同的張量。 否則，將返回一個新的 tensordict，其結構與 self 相同，但特徵維度已縮減。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• dtype (torch.dtype, optional) – 返回張量的所需資料型別。如果指定，則在執行操作之前將輸入張量強制轉換為 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() → T

計算 TensorDict 中每個元素的 tan() 值。

tan_() → T

原地計算 TensorDict 中每個元素的 tan() 值。

tanh() → T

計算 TensorDict 中每個元素的 tanh() 值。

tanh_() → T

原地計算 TensorDict 中每個元素的 tanh() 值。

to(*args, **kwargs: Any) → PersistentTensorDict

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

不允許將張量轉換為新的 dtype，因為 tensordict 不一定包含單一的張量 dtype。

引數:

• device (torch.device, optional) – tensordict 的目標裝置。

• dtype (torch.dtype, optional) – tensordict 的目標浮點或複數 dtype。

• tensor (torch.Tensor, optional) – dtype 和裝置是此 TensorDict 中所有張量的目標 dtype 和裝置的張量。

關鍵字引數:

• non_blocking (bool, optional) – 操作是否應為阻塞。

• memory_format (torch.memory_format, optional) – 此 tensordict 中 4D 引數和緩衝區的目標記憶體格式。

• batch_size (torch.Size, optional) – 輸出 tensordict 的結果批大小。

• other (TensorDictBase, optional) –

  TensorDict 例項，其 dtype 和裝置是此 TensorDict 中所有張量的目標 dtype 和裝置。

  注意

  由於 TensorDictBase 例項沒有 dtype，因此 dtype 從示例葉子節點收集。 如果有多個 dtype，則不進行 dtype 轉換。

• non_blocking_pin (bool, optional) –

  如果 True，則張量在傳送到裝置之前被固定。 這將非同步完成，但可以透過 num_threads 引數控制。

  注意

  呼叫 tensordict.pin_memory().to("cuda") 通常比 tensordict.to("cuda", non_blocking_pin=True) 慢得多，因為 pin_memory 在第二種情況下是非同步呼叫的。 如果張量很大且數量眾多，則多執行緒 pin_memory 通常是有益的：當要傳送的張量太少時，生成執行緒和收集資料的開銷會超過多執行緒的優勢，並且如果張量很小，則迭代長列表的開銷也非常大。

• num_threads (int or None, optional) – 如果 non_blocking_pin=True，則用於 pin_memory 的執行緒數。 預設情況下，將生成 max(1, torch.get_num_threads()) 個執行緒。 num_threads=0 將取消 pin_memory() 呼叫的任何多執行緒。

• inplace (bool, optional) – 如果 True，則資料將就地寫入同一 tensordict 中。 如果構建 tensordict 受 CPU 開銷限制，這可能會顯著加快速度。 預設為 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) → dict[str, Any]

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

引數:

retain_none (bool) – 如果 True，則 tensorclass 例項中的 None 值將寫入字典中。 否則，它們將被丟棄。 預設值：True。

to_h5(filename, **kwargs)

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

引數:

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

• **kwargs – 要傳遞給 h5py.File.create_dataset() 的 kwargs。

返回值:

一個連結到新建立的檔案的 PersitentTensorDict 例項。

示例

>>> import tempfile
>>> import timeit
>>>
>>> from tensordict import TensorDict, MemoryMappedTensor
>>> td = TensorDict({
...     "a": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000)),
...     "b": {"c": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000, 3))},
... }, [1_000_000])
>>>
>>> file = tempfile.NamedTemporaryFile()
>>> td_h5 = td.to_h5(file.name, compression="gzip", compression_opts=9)
>>> print(td_h5)
PersistentTensorDict(
    fields={
        a: Tensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float32, is_shared=False),
        b: PersistentTensorDict(
            fields={
                c: Tensor(shape=torch.Size([1000000, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([1000000]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([1000000]),
    device=None,
    is_shared=False)

to_module(module: Module, *, inplace: bool | None = None, return_swap: bool = True, swap_dest=None, use_state_dict: bool = False, non_blocking: bool = False, memo=None)

將 TensorDictBase 例項的內容遞迴地寫入給定的 nn.Module 屬性。

to_module 也可以用作上下文管理器，以臨時使用引數/緩衝區集合填充模組（參見下面的示例）。

引數:

module (nn.Module) – 要將引數寫入的模組。

關鍵字引數:

• inplace (bool, optional) – 如果為 True，則模組中的引數或張量將就地更新。預設為 False。

• return_swap (bool, optional) – 如果為 True，將返回舊的引數配置。預設為 False。

• swap_dest (TensorDictBase, optional) – 如果 return_swap 為 True，則應寫入交換的 tensordict。

• use_state_dict (bool, optional) – 如果為 True，將使用 state-dict API 載入引數（包括 state-dict 鉤子）。預設為 False。

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

示例

>>> from torch import nn
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = TensorDict.from_module(module)
>>> params.data.zero_()
>>> params.to_module(module)
>>> assert (module.layers[0].linear1.weight == 0).all()

使用 tensordict 作為上下文管理器對於進行函式式呼叫非常有用： .. rubric:: 示例

>>> from tensordict import from_module
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = TensorDict.from_module(module)
>>> params = params.data * 0 # Use TensorDictParams to remake these tensors regular nn.Parameter instances
>>> with params.to_module(module):
...     # Call the module with zeroed params
...     y = module(*inputs)
>>> # The module is repopulated with its original params
>>> assert (TensorDict.from_module(module) != 0).any()

返回值:

如果 return_swap 為 True，則為包含模組值的 tensordict，否則為 None。

to_namedtuple(dest_cls: Optional[type] = None)

將 tensordict 轉換為 namedtuple。

引數:

dest_cls (Type, optional) – 可選的 namedtuple 類以供使用。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> data.to_namedtuple()
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))

to_padded_tensor(padding=0.0, mask_key: Optional[NestedKey] = None) → T

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

引數:

• padding (float) – tensordict 中張量的填充值。預設為 0.0。

• mask_key (NestedKey, optional) – 如果提供，則將寫入有效值的掩碼的鍵。 如果異構維度不是 tensordict 批次大小的一部分，將導致錯誤。 預設為 None

to_pytree()

將 tensordict 轉換為 PyTree。

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

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

to_struct_array()

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

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

參見

from_struct_array() 瞭解更多資訊。

返回值:

輸入 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: Optional[bool] = None) → T

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

引數:

retain_none (bool) –

如果為 True，則來自 tensorclass 例項的 None 值將被寫入 tensordict。 否則，它們將被丟棄。 預設值：True。

注意

從 v0.8 開始，預設值將切換為 False。

返回值:

包含相同值的新 TensorDict 物件。

transpose(dim0, dim1)

返回一個 tensordict，它是輸入的轉置版本。給定的維度 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() → T

計算 TensorDict 的每個元素的 trunc() 值。

trunc_() → T

就地計算 TensorDict 的每個元素的 trunc() 值。

type(dst_type)

將所有張量轉換為 dst_type。

引數:

dst_type (type 或 string) – 所需的型別

uint16()

將所有 tensors 轉換為 torch.uint16。

uint32()

將所有 tensors 轉換為 torch.uint32。

uint64()

將所有 tensors 轉換為 torch.uint64。

uint8()

將所有 tensors 轉換為 torch.uint8。

unbind(dim: int) → tuple[T, ...]

返回沿指定維度解綁的索引化 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) – 指定要展開的輸入 tensor 的維度。

• unflattened_size (shape) – 是 tensordict 的未展開維度的新形狀。

示例

>>> td = TensorDict({
...     "a": torch.arange(60).view(3, 4, 5),
...     "b": torch.arange(12).view(3, 4)},
...     batch_size=[3, 4])
>>> td_flat = td.flatten(0, 1)
>>> td_unflat = td_flat.unflatten(0, [3, 4])
>>> assert (td == td_unflat).all()

unflatten_keys(separator: str = '.', inplace: bool = False) → T

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

TensorDict 型別將會丟失，結果將是一個簡單的 TensorDict 例項。巢狀 tensordict 的元資料將從根推斷：資料樹中的所有例項將共享相同的批次大小、維度名稱和裝置。

引數:

• separator (str, 可選) – 巢狀項之間的分隔符。

• inplace (bool, 可選) – 如果 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() 在處理狀態字典時特別有用，因為它們可以無縫地將扁平的字典轉換為模仿模型結構的資料結構。

示例

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

unlock_() → T

解鎖一個 tensordict 以進行非原地操作。

可以用作裝飾器。

參見 lock_() 瞭解更多詳情。

unsqueeze(*args, **kwargs)

對所有張量在 -td.batch_dims 和 td.batch_dims 之間的維度進行unsqueeze操作，並將它們返回到一個新的 tensordict 中。

引數:

dim (int) – 要沿其進行unsqueeze操作的維度

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> td = td.unsqueeze(-2)
>>> td.shape
torch.Size([3, 1, 4])
>>> td.get("x").shape
torch.Size([3, 1, 4, 2])

該操作也可以用作上下文管理器。對原始 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: Union[dict[str, torch.Tensor], T], clone: bool = False, inplace: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None, is_leaf: Optional[Callable[[Type], bool]] = None) → T

使用字典或另一個 TensorDict 的值來更新 TensorDict。

引數:

• input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。

• clone (bool, optional) – 輸入的 (tensor) 字典中的張量是否應在設定之前被克隆。 預設為 False。

• inplace (bool, optional) – 如果為 True 並且鍵與 tensordict 中已存在的鍵匹配，則該鍵值對的更新將以原地方式發生。 如果找不到該條目，則會被新增。預設為 False。

關鍵字引數:

• keys_to_update (NestedKeys 序列, 可選) – 如果提供，則僅更新 key_to_update 中的鍵列表。 旨在避免呼叫 data_dest.update(data_src.select(*keys_to_update))。

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

• is_leaf (Callable[[Type], bool], 可選) –

  一個可呼叫物件，指示一個物件型別是否被認為是葉子並被交換，還是一個張量集合。

  參見

  is_leaf_nontensor() 和 default_is_leaf()。

返回值:

self

示例

>>> td = TensorDict({}, batch_size=[3])
>>> a = torch.randn(3)
>>> b = torch.randn(3, 4)
>>> other_td = TensorDict({"a": a, "b": b}, batch_size=[])
>>> td.update(other_td, inplace=True) # writes "a" and "b" even though they can't be found
>>> assert td['a'] is other_td['a']
>>> other_td = other_td.clone().zero_()
>>> td.update(other_td)
>>> assert td['a'] is not other_td['a']

update_(input_dict_or_td: Union[dict[str, torch.Tensor], T], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) → T

使用字典或另一個 TensorDict 中的值就地更新 TensorDict。

與 update() 不同，如果 self 不知道該鍵，此函式將丟擲錯誤。

引數:

• input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。

• clone (bool, optional) – 輸入的 (tensor) 字典中的張量是否應在設定之前被克隆。 預設為 False。

關鍵字引數:

• keys_to_update (NestedKeys 序列, 可選) – 如果提供，則僅更新 key_to_update 中的鍵列表。 它的目的是避免呼叫 data_dest.update_(data_src.select(*keys_to_update))。

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

返回值:

self

示例

>>> a = torch.randn(3)
>>> b = torch.randn(3, 4)
>>> td = TensorDict({"a": a, "b": b}, batch_size=[3])
>>> other_td = TensorDict({"a": a*0, "b": b*0}, batch_size=[])
>>> td.update_(other_td)
>>> assert td['a'] is not other_td['a']
>>> assert (td['a'] == other_td['a']).all()
>>> assert (td['a'] == 0).all()

update_at_(input_dict_or_td: Union[dict[str, torch.Tensor], T], idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) → T

使用字典或另一個 TensorDict 中的值，在指定的索引處就地更新 TensorDict。

與 TensorDict.update 不同，如果 TensorDict 不知道該鍵，此函式將丟擲錯誤。

引數:

• input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。

• idx (int, torch.Tensor, iterable, slice) – 應該發生更新的 tensordict 的索引。

• clone (bool, optional) – 是否應該在設定之前克隆輸入 ( tensor) 字典中的張量。 預設為 False。

關鍵字引數:

• keys_to_update (NestedKeys 序列, 可選) – 如果提供，則僅更新 key_to_update 中的鍵列表。

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

返回值:

self

示例

>>> td = TensorDict({
...     'a': torch.zeros(3, 4, 5),
...     'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4])
>>> td.update_at_(
...     TensorDict({
...         'a': torch.ones(1, 4, 5),
...         'b': torch.ones(1, 4, 10)}, batch_size=[1, 4]),
...    slice(1, 2))
TensorDict(
    fields={
        a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32),
        b: Tensor(torch.Size([3, 4, 10]), dtype=torch.float32)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> assert (td[1] == 1).all()

values(include_nested: bool = False, leaves_only: bool = False, is_leaf=None, *, sort: bool = False) → Iterator[Tensor]

返回一個表示 tensordict 值的生成器。

引數:

• include_nested (bool, optional) – 如果為 True, 將返回巢狀值。 預設為 False。

• leaves_only (bool, optional) – 如果為 False, 將僅返回葉節點。 預設為 False。

• is_leaf (callable, optional) –

  一個可呼叫物件，作用於一個類型別，返回一個布林值，指示是否應將該類視為葉子節點。

  注意

  is_leaf 的目的不是阻止遞迴呼叫到巢狀的 tensordict 中，而是當 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) → tensordict.base.TensorDictBase | torch.Tensor

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

引數:

• dim (int, int 元組, 可選) – 如果 None，則返回一個無維度的 tensordict，其中包含所有葉子的總和值（如果可以計算）。如果是整數或整數元組，則僅當此維度與 tensordict 形狀相容時，才對指定的維度呼叫 var。目前只允許使用 “feature” 字串。使用 dim=”feature” 將實現對所有特徵維度的縮減。如果 reduce=True，則將返回一個形狀為 TensorDict 批大小的張量。否則，將返回一個新的 tensordict，其結構與 self 相同，但特徵維度已縮減。

• keepdim (bool) – 輸出張量是否保留維度。

關鍵字引數:

• correction (int) – 樣本大小和樣本自由度之間的差異。預設為貝塞爾校正，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 批大小相容的新形狀的張量的檢視。

或者，可以提供 dtype 作為第一個未命名的引數。在這種情況下，所有張量將使用相應的 dtype 進行檢視。請注意，這假定新形狀將與提供的 dtype 相容。有關 dtype 檢視的更多資訊，請參閱 view()。

引數:

• *shape (int) – 結果 tensordict 的新形狀。

• dtype (torch.dtype) – 或者，用於表示張量內容的 dtype。

• size – 可迭代物件

關鍵字引數:

batch_size (torch.Size, 可選) – 如果提供了 dtype，則可以使用此關鍵字引數重置批大小。 如果使用形狀呼叫 view，則這不起作用。

返回值:

具有所需批大小的新 tensordict。

示例

>>> td = TensorDict(source={'a': torch.zeros(3,4,5),
...    'b': torch.zeros(3,4,10,1)}, batch_size=torch.Size([3, 4]))
>>> td_view = td.view(12)
>>> print(td_view.get("a").shape)  # torch.Size([12, 5])
>>> print(td_view.get("b").shape)  # torch.Size([12, 10, 1])
>>> td_view = td.view(-1, 4, 3)
>>> print(td_view.get("a").shape)  # torch.Size([1, 4, 3, 5])
>>> print(td_view.get("b").shape)  # torch.Size([1, 4, 3, 10, 1])

where(condition, other, *, out=None, pad=None)

返回一個 TensorDict，其中包含從 self 或 other 中選擇的元素，具體取決於 condition。

引數:

• condition (BoolTensor) – 當 True (非零) 時，產生 self，否則產生 other。

• other (TensorDictBase 或 標量) – 在 condition 為 False 的索引處選擇的值（如果 other 是標量）。

關鍵字引數:

• out (TensorDictBase, 可選) – 輸出的 TensorDictBase 例項。

• pad (標量, 可選) – 如果提供，則源或目標 tensordict 中缺少的鍵將被寫入為 torch.where(mask, self, pad) 或 torch.where(mask, pad, other)。 預設為 None，即不容忍缺少鍵。

zero_() → T

就地將 tensordict 中的所有張量歸零。

zero_grad(set_to_none: bool = True) → T

遞迴地將 TensorDict 的所有梯度歸零。

引數:

set_to_none (bool, 可選) – 如果 True，則 tensor.grad 將為 None，否則為 0。預設為 True。