快捷方式

NonTensorStack

class tensordict.NonTensorStack(*args, **kwargs)

一個簡單的包裝器,圍繞 LazyStackedTensorDict,使非張量資料的堆疊更容易識別。

當在 NonTensorDataNonTensorStack 的列表上呼叫 stack() 時,將返回一個 NonTensorStack

示例

>>> from tensordict import NonTensorData
>>> import torch
>>> data = torch.stack([
...     torch.stack([NonTensorData(data=(i, j), batch_size=[]) for i in range(2)])
...    for j in range(3)])
>>> print(data)
NonTensorStack(
    [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(0, 2), (1, ...,
    batch_size=torch.Size([3, 2]),
    device=None)

要獲取儲存在 NonTensorStack 中的值,請呼叫 tolist

abs() Any

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

abs_() Any

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

acos() Any

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

acos_() Any

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

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

other 乘以 alpha 並加到 self 上。

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

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

關鍵字引數:
  • alpha (數字, 可選) – other 的乘數。

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

add_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor | float, *, alpha: float | None = None) Any

的就地版本add()

注意

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

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

執行 other1 除以 other2 的逐元素除法,將結果乘以標量 value 並加到 self 上。

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

selfother1other2 的形狀必須可廣播。

對於 FloatTensorDoubleTensor 型別的輸入,value 必須是實數,否則為整數。

引數:
  • other1 (TensorDictTensor) – 被除數 tensordict(或 tensor)

  • tensor2 (TensorDictTensor) – 除數 tensordict(或 tensor)

關鍵字引數:

value (數字, 可選) – \(\text{tensor1} / \text{tensor2}\) 的乘數

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

的就地版本addcdiv()

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

執行 other1other2 的逐元素乘法,將結果乘以標量 value 並加到 self 上。

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

selfother1other2 的形狀必須可廣播。

對於 FloatTensorDoubleTensor 型別的輸入,value 必須是實數,否則為整數。

引數:
  • other1 (TensorDictTensor) – 要相乘的 tensordict 或 tensor

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

關鍵字引數:

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

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

的就地版本addcmul()

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

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

引數:

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

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

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

max() 相同,但 return_indices=False

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

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

min() 相同,但 return_indices=False

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

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

引數:

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

append(tensordict: T) None

將 TensorDict 追加到堆疊。

類似於 list.append。追加的 TensorDict 必須具有相容的 batch_size 和 device。追加操作是就地進行的,不返回任何內容。

引數:

tensordict (TensorDictBase) – 要追加到堆疊中的 TensorDict。

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[Any]

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

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

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

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

關鍵字引數:
  • batch_size (sequence of int, optional) – 如果提供,結果 TensorDict 將具有所需的 batch_size。batch_size 引數應與轉換後的 batch_size 匹配。這是一個僅限關鍵字的引數。

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

  • names (字串列表, 可選) – 新的維度名稱,以防batch_size被修改。

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

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

  • filter_empty (bool, optional) – 如果為 True,則將過濾掉空的 tensordict。這也帶來了較低的計算成本,因為不會建立和銷燬空的 Datal。非張量資料被視為葉子,因此即使函式未觸及,它也會保留在 tensordict 中。預設為 False 以保持向後相容。

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

  • call_on_nested (bool, optional) –

    如果為 True,則該函式將應用於第一級張量和容器(TensorDict 或 tensorclass)。在這種情況下,func 負責傳播其呼叫到巢狀級別。這允許在將呼叫傳播到巢狀 tensordicts 時進行精細控制。如果為 False,則該函式僅應用於葉子節點,並且 apply 將負責將函式分派到所有葉子節點。

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

  • out (TensorDictBase, 可選) –

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

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

    警告

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

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

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

返回:

一個包含轉換後張量的新tensordict。

示例

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

注意

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

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

注意

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

>>> out = td.clone(False).update(td.apply(...))
apply_(fn: Callable, *others, **kwargs)

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

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

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

關鍵字引數:請參閱 apply()

返回:

self 或 self 的副本,其中應用了函式

as_tensor() Any

將 tensordict 的每個葉子節點轉換為普通的 torch.Tensor。

asin() Any

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

asin_() Any

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

atan() Any

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

atan_() Any

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

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

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

引數:
  • batch_dims (int, 可選) – 如果提供,則批次大小最多為 batch_dims

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

返回:

self

示例

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

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

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

property batch_dims: int

tensordict批次大小的長度。

返回:

描述tensordict維度的整數。

property batch_size: Size

TensorDict的形狀(或batch_size)。

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

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

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

返回:

一個 Size 物件,描述 TensorDict 的批次大小。

示例

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

將所有張量轉換為torch.bfloat16

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

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

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

other (TensorDictBasetorch.Tensor) – 要執行按位 AND 的張量或 TensorDict。

關鍵字引數:

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

bool() Any

將所有張量轉換為torch.bool

bytes(*, count_duplicates: bool = True) int

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

關鍵字引數:

count_duplicates (bool) – 是否計算重複的張量作為獨立的張量。如果為 False,則只丟棄完全相同的張量(來自同一基礎張量的相同檢視但具有不同 ID 的張量將被計算兩次)。預設為 True(每個張量假定為單個副本)。

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

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

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

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

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

引數:

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

關鍵字引數:
  • sorted (boollist of NestedKeys) – 如果為 True,則條目將按字母順序連線。如果為 False(預設),則使用字典順序。或者,可以提供一個鍵名列表,並相應地連線張量。這會產生一些開銷,因為鍵列表需要與 tensordict 中的葉子名稱列表進行檢查。

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

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

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

引數:

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

關鍵字引數:
  • out_key (NestedKey) – 連線輸入的新鍵名。

  • keep_entries (bool, optional) – 如果為 False,則 keys 中的條目將被刪除。預設為 False

  • dim (int, optional) – 必須進行連線的維度。預設為 0

返回: self

示例

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

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

ceil_() Any

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

chunk(chunks: int, dim: int = 0) tuple[tensordict._tensorcollection.TensorCollection, ...]

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

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

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

  • dim (int, 可選) – 分割tensordict的維度。預設為0。

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> td0, td1 = td.chunk(dim=-1, chunks=2)
>>> td0['x']
tensor([[[ 0,  1],
         [ 2,  3]],
        [[ 8,  9],
         [10, 11]],
        [[16, 17],
         [18, 19]]])
clamp(min: tensordict.base.TensorDictBase | torch.Tensor = None, max: tensordict.base.TensorDictBase | torch.Tensor = None, *, out=None) Any

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

令 min_value 和 max_value 分別為 minmax,這返回

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

如果 minNone,則沒有下限。或者,如果 maxNone,則沒有上限。

注意

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

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

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

引數:

other (TensorDictTensor) – 另一個輸入tensordict或張量。

關鍵字引數:

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

clamp_max_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) Any

的就地版本clamp_max()

注意

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

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

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

引數:

other (TensorDictTensor) – 另一個輸入tensordict或張量。

關鍵字引數:

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

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

的就地版本clamp_min()

注意

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

clear() Any

清除 tensordict 的內容。

clear_device_() Any

清除 tensordict 的裝置。

返回: self

clear_refs_for_compile_() Any

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

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

返回: self

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

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

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

引數:

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

注意

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

complex128() Any

將所有張量轉換為 torch.complex128

complex32() Any

將所有張量轉換為 torch.complex32

complex64() Any

將所有張量轉換為 torch.complex64

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

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

引數:

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

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

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

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

  • inplace (bool, optional) – 如果為 True,則結果 tensordict 與 self 相同,值為更新。預設為 False

  • return_early (bool, optional) – 如果為 Truenum_threads>0,則該方法將返回一個 tensordict 的 future。可以透過 future.result() 查詢結果 tensordict。

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

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

  • pin_memory (bool, optional) – 整合資料是否應放在固定記憶體中。預設為 False

  • metadata (bool, optional) – 如果為 True,則元資料將與公共儲存一起儲存。如果提供了檔名,則此選項無效。儲存元資料對於控制序列化如何實現可能很有用,因為 TensorDict 會以不同的方式處理已整合 TD 的 pickling/unpickling,無論元資料是否可用。

注意

如果 tensordict 已經整合,則所有引數將被忽略,並返回 self。呼叫 contiguous() 以重新整合。

示例

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

返回一個相同型別的新 tensordict,其值是連續的(如果值已經是連續的,則返回 self)。(Returns a new tensordict of the same type with contiguous values (or self if values are already contiguous).)

copy() Any

返回 tensordict 的淺複製(即,複製結構但不復制資料)。

等同於 TensorDictBase.clone(recurse=False)

copy_(tensordict: T, non_blocking: bool = False) Any

請參閱 TensorDictBase.update_

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

copy_at_(tensordict: T, idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], non_blocking: bool = False) Any

請參閱 TensorDictBase.update_at_

cos() Any

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

cos_() Any

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

cosh() Any

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

cosh_() Any

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

cpu(**kwargs) Any

將 tensordict 轉換為 CPU。

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

create_nested(key)

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

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

示例

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

將 tensordict 轉換為 cuda 裝置(如果尚未轉換)。

引數:

device (int, optional) – 如果提供,則為張量應被轉換到的 cuda 裝置。

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

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

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

引數:

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

關鍵字引數:
  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

  • return_argmins (bool, optional) – 當 dim 引數被傳遞時,cummax() 會返回一個包含值和索引的命名元組。這個 TensorDict 的等價操作是返回一個包含 "values""indices" 條目的 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.cummax(dim=0)
cummax(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummax(reduce=True, dim=0)
torch.return_types.cummax(...)
cummin(dim: int, *, reduce: Optional[bool] = None, return_indices: bool = True) Union[Any, Tensor]

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

引數:

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

關鍵字引數:
  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

  • return_argmins (bool, optional) – 當 dim 引數被傳遞時,cummin() 會返回一個包含值和索引的命名元組。這個 TensorDict 的等價操作是返回一個包含 "values""indices" 條目的 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.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

嘗試返回堆疊中的唯一值。

如果存在多個唯一值,則引發 ValueError。

data_ptr(*, storage: bool = False)

返回tensordict葉子節點的data_ptr。

這有助於檢查兩個tensordict是否共享相同的data_ptr()

關鍵字引數:

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

示例

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

注意

LazyStackedTensorDict 例項將顯示為巢狀的 tensordicts,以反映其葉節點的實際 data_ptr()

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

刪除tensordict的鍵。

引數:

key (NestedKey) – 要刪除的鍵

返回:

self

densify(layout: layout = torch.strided)

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

關鍵字引數:

layout (torch.layout) – 巢狀張量的佈局,如果有的話。預設為 strided

property depth: int

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

最小深度為0(沒有巢狀的tensordict)。

detach() Any

分離tensordict中的張量。

返回:

返回一個不包含需要梯度的張量的新tensordict。

detach_() Any

就地分離tensordict中的張量。

返回:

self。

property device: torch.device | None

TensorDict的裝置。

如果 TensorDict 指定了裝置,則其所有張量(包括巢狀張量)必須位於同一裝置上。如果 TensorDict 裝置為 None,則不同值可以位於不同裝置上。

返回:

torch.device 物件,指示張量所在的位置,如果TensorDict沒有裝置則為None。

示例

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

請參閱 batch_dims()

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

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

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

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

引數:

other (TensorDict, TensorNumber) – 除數。

關鍵字引數:

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

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

原地版本的 div()

注意

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

double() Any

將所有張量轉換為torch.bool

property dtype

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

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

將tensordict儲存到磁碟。

此函式是 memmap() 的代理。

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

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

引數:

recurse (bool, optional) – 如果為 True,則將複製 TensorDict 的整個結構而沒有內容。否則,僅複製根。預設為 False

關鍵字引數:
  • batch_size (torch.Size, optional) – tensordict 的新批次大小。

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

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

entry_class(key: NestedKey) type

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

get() 執行可能很昂貴時,此方法應優先於 tensordict.get(key).shape

erf() Any

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

erf_() Any

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

erfc() Any

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

erfc_() Any

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

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

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

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

引數:
  • *keys (str) – 要排除的鍵。

  • inplace (bool) – 如果為 True,則 tensordict 將就地剪枝。預設為 False

返回:

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

示例

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

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

exp_() Any

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

expand(*args: int, inplace: bool = False) Any

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

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

示例

>>> td = TensorDict({
...     'a': torch.zeros(3, 4, 5),
...     'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4])
>>> td_expand = td.expand(10, 3, 4)
>>> assert td_expand.shape == torch.Size([10, 3, 4])
>>> assert td_expand.get("a").shape == torch.Size([10, 3, 4, 5])
expand_as(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) Any

將tensordict的形狀廣播到other的形狀,並相應地擴充套件它。

如果輸入是張量集合(tensordict或tensorclass),則葉子節點將進行一對一的擴充套件。

示例

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

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

expm1_() Any

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

extend(tensordict: Union[list[T], T]) None

使用新的 tensordict 擴充套件惰性堆疊。

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

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

引數:
  • key (str or nested key) – entry to be filled.

  • value (Number or bool) – value to use for the filling.

返回:

self

filter_empty_()

就地過濾掉所有空的tensordict。

filter_non_tensor_data() Any

過濾掉所有非張量資料。

flatten(start_dim: int = 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, is_leaf: Optional[Callable[[Type], bool]] = None) Any

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

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

引數:
  • separator (str, optional) – 巢狀項之間的分隔符。

  • inplace (bool, optional) – 如果為 True,則結果 tensordict 將與呼叫它的 tensordict 具有相同的身份。預設為 False

  • is_leaf (callable, optional) –

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

    注意

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

    另請參閱

    is_leaf_nontensor()default_is_leaf()

示例

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

此方法和 unflatten_keys() 方法在處理 state-dict 時非常有用,因為它們使得能夠無縫地將扁平化字典轉換為模仿模型結構的結構。

示例

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

將所有張量轉換為torch.float

float16() Any

將所有張量轉換為torch.float16

float32() Any

將所有張量轉換為torch.float32

float64() Any

將所有張量轉換為torch.float64

floor() Any

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

floor_() Any

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

frac() Any

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

frac_() Any

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

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

Recursively converts any object to a TensorDict.

注意

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

注意

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

引數:

obj – The object to be converted.

關鍵字引數:
  • auto_batch_size (bool, optional) – 如果 True,則自動計算並應用批次大小。預設為 False

  • batch_dims (int, optional) – 如果 auto_batch_sizeTrue,則定義輸出 tensordict 應具有的維度數量。預設為 None(每個級別都有完整的批次大小)。

  • device (torch.device, optional) – 將建立 TensorDict 的裝置。

  • batch_size (torch.Size, optional) – TensorDict 的批次大小。與 auto_batch_size 互斥。

返回:

A TensorDict representation of the input object.

Supported objects

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

Converts a dataclass into a TensorDict instance.

引數:

dataclass – The dataclass instance to be converted.

關鍵字引數:
  • dest_cls (tensorclass, optional) – 用於對映資料的 tensorclass 型別。如果未提供,則建立一個新類。如果 `obj` 是一個型別或 as_tensorclass 為 `False`,則無效。

  • auto_batch_size (bool, optional) – 如果 True,將自動確定並應用批次大小到生成的 TensorDict。預設為 False

  • batch_dims (int, optional) – 如果 auto_batch_sizeTrue,則定義輸出 tensordict 應具有的維度數量。預設為 None(每個級別都有完整的批次大小)。

  • as_tensorclass (bool, optional) – 如果 True,則委託給自由函式 from_dataclass() 進行轉換,並返回一個 tensor 相容的類(tensorclass())或例項,而不是 TensorDict。預設為 False

  • device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None

  • batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None

返回:

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

丟擲:

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

警告

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

注意

  • 此方法建立一個新的 TensorDict 例項,其鍵對應於輸入 dataclass 的欄位。

  • 結果 TensorDict 中的每個鍵都使用 `cls.from_any` 方法進行初始化。

  • auto_batch_size 選項允許自動確定批次大小並將其應用於結果 TensorDict。

classmethod from_dict(input_dict: Union[Dict[str, Dict[NestedKey, Any]], List[Dict[NestedKey, Any]]], *other, auto_batch_size: bool = False, batch_size=None, device=None, batch_dims=None, stack_dim_name=None, stack_dim=0, names: Optional[list[str]] = None)

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

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

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

引數:

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

關鍵字引數:
  • auto_batch_size (bool, optional) – 如果 True,則自動計算並應用批次大小。預設為 False

  • batch_size (iterable of int, optional) – a batch size for the tensordict.

  • device (torch.device相容型別, optional) – TensorDict 的裝置。

  • batch_dims (int, optional) – batch_dims(即用於 batch_size 的前導維度數量)。與 batch_size 互斥。請注意,這是 tensordict 的 __最大__ 批次維度數,較小的數字也可以接受。

  • names (list of str, optional) – the dimension names of the tensordict.

示例

>>> input_dict = {"a": torch.randn(3, 4), "b": torch.randn(3)}
>>> print(TensorDict.from_dict(input_dict))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> # nested dict: the nested TensorDict can have a different batch-size
>>> # as long as its leading dims match.
>>> input_dict = {"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}
>>> print(TensorDict.from_dict(input_dict))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> # we can also use this to work out the batch sie of a tensordict
>>> input_td = TensorDict({"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}, [])
>>> print(TensorDict.from_dict(input_td))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
from_dict_instance(input_dict, *, auto_batch_size: 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: str = 'r', auto_batch_size: bool = False, batch_dims: Optional[int] = None, batch_size: Optional[Size] = None)

從 h5 檔案建立 PersistentTensorDict。

引數:

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

關鍵字引數:
  • mode (str, 可選) – 讀取模式。預設為 "r"

  • auto_batch_size (bool, optional) – 如果 True,則自動計算批次大小。預設為 False

  • batch_dims (int, optional) – 如果 auto_batch_sizeTrue,則定義輸出 tensordict 應具有的維度數量。預設為 None(每個級別都有完整的批次大小)。

  • batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None

返回:

輸入 h5 檔案的 PersistentTensorDict 表示。

示例

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

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

引數:
  • module (nn.Module) – 要從中獲取引數的模組。

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

  • lock (bool, optional) – 如果為 True,則結果 tensordict 將被鎖定。預設為 True

  • use_state_dict (bool, optional) –

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

    注意

    這在使用 state-dict hook 時尤其有用。

示例

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

為 vmap 的 ensemable 學習/特徵期望應用檢索多個模組的引數。

引數:

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

關鍵字引數:
  • as_module (bool, optional) – 如果 True,將返回一個 TensorDictParams 例項,該例項可用於在 torch.nn.Module 中儲存引數。預設為 False

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

  • use_state_dict (bool, optional) –

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

    注意

    這在使用 state-dict hook 時尤其有用。

  • lazy_stack (bool, optional) –

    是否密集堆疊或懶惰堆疊引數。預設為 False(密集堆疊)。

    注意

    lazy_stackas_module 是互斥的特性。

    警告

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

    警告

    雖然使用惰性堆疊來保留原始引數引用可能很誘人,但請記住,惰性堆疊每次呼叫 get() 時都會執行堆疊操作。這將需要記憶體(N 倍的引數大小,如果構建了圖,則更多)和時間來計算。它還意味著最佳化器將包含更多引數,並且像 step()zero_grad() 這樣的操作將花費更長的時間來執行。總的來說,lazy_stack 應保留用於極少數用例。

  • expand_identical (bool, optional) – 如果 True 且相同的引數(相同的標識)被堆疊到自身,則將返回該引數的擴充套件版本。當 lazy_stack=True 時,將忽略此引數。

示例

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

lazy_stack=True 時,情況略有不同

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

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

引數:

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

關鍵字引數:
  • auto_batch_size (bool, optional) – 如果 True,則自動計算並應用批次大小。預設為 False

  • batch_dims (int, optional) – 如果 auto_batch_sizeTrue,則定義輸出 tensordict 應具有的維度數量。預設為 None(每個級別都有完整的批次大小)。

  • device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None

  • batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None

返回:

輸入命名元組的 TensorDict 表示。

示例

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

將 pytree 轉換為 TensorDict 例項。

此方法旨在儘可能保留 pytree 的巢狀結構。

其他非張量鍵將被新增,以跟蹤每個級別的標識,從而提供內建的 pytree 到 tensordict 的雙射轉換 API。

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

注意

對於字典,非 `NestedKey` 鍵被單獨註冊為 `NonTensorData` 例項。

注意

可轉換為張量型別(如 int、float 或 np.ndarray)將被轉換為 torch.Tensor 例項。請注意,此轉換是滿射的:將 tensordict 轉換回 pytree 將無法恢復原始型別。

示例

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

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

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

引數:
  • src (int) – 傳送元資料的源程序的秩。

  • group ("ProcessGroup", optional) – 要使用的程序組。預設為 None。

  • device (torch.device, 可選) – 用於張量運算的裝置。預設為 None。

返回:

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

返回型別:

TensorDict

另請參閱

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

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

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

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

注意

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

引數:

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

關鍵字引數:
  • auto_batch_size (bool, optional) – 如果 True,則自動計算批次大小。預設為 False

  • batch_dims (int, optional) – 如果 auto_batch_sizeTrue,則定義輸出 tensordict 應具有的維度數量。預設為 None(每個級別都有完整的批次大小)。

  • device (torch.device, 可選) –

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

    注意

    更改裝置(即,指定任何非 `None` 或 `"cpu"` 的裝置)將傳輸資料,從而導致返回資料的記憶體位置發生更改。

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

返回:

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

示例

>>> x = np.array(
...     [("Rex", 9, 81.0), ("Fido", 3, 27.0)],
...     dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")],
... )
>>> td = TensorDict.from_struct_array(x)
>>> x_recon = td.to_struct_array()
>>> assert (x_recon == x).all()
>>> assert x_recon.shape == x.shape
>>> # Try modifying x age field and check effect on td
>>> x["age"] += 1
>>> assert (td["age"] == np.array([10, 4])).all()
classmethod from_tuple(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)

將元組轉換為 TensorDict。

引數:

obj – 要轉換的元組例項。

關鍵字引數:
  • auto_batch_size (bool, optional) – 如果 True,則自動計算批次大小。預設為 False

  • batch_dims (int, optional) – 如果 auto_batch_sizeTrue,則定義輸出 tensordict 應具有的維度數量。預設為 None(每個級別都有完整的批次大小)。

  • device (torch.device, optional) – 將建立 TensorDict 的裝置。預設為 None

  • batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None

返回:

輸入的元組的 TensorDict 表示。

示例

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

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

引數:
  • keys (list of NestedKey) – 指定新字典鍵的可迭代物件。

  • value (compatible type, optional) – 所有鍵的值。預設為 0

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

沿由 dim 指定的軸收集值。

引數:
  • dim (int) – 要收集元素的維度

  • index (torch.Tensor) – 一個長整型張量,其維度數量與 tensordict 匹配,只有一個維度不同(收集維度)。它的元素引用沿所需維度收集的索引。

  • out (TensorDictBase, 可選) – 目標 tensordict。它必須與索引具有相同的形狀。

示例

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

Gather 保留維度名稱。

示例

>>> td.names = ["a", "b"]
>>> td_gather = td.gather(dim=1, index=index)
>>> td_gather.names
["a", "b"]
gather_and_stack(dst: int, group: 'torch.distributed.ProcessGroup' | None = None) Self | None

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

引數:
  • dst (int) – 目標工作程序的 rank,其中將呼叫 gather_and_stack()

  • group (torch.distributed.ProcessGroup, 可選) – 如果設定,將使用指定的程序組進行通訊。否則,將使用預設程序組。預設為 `None`。

示例

>>> from torch import multiprocessing as mp
>>> from tensordict import TensorDict
>>> import torch
>>>
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://:10003",
...     )
...     # Create a single tensordict to be sent to server
...     td = TensorDict(
...         {("a", "b"): torch.randn(2),
...          "c": torch.randn(2)}, [2]
...     )
...     td.gather_and_stack(0)
...
>>> def server():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=0,
...         world_size=2,
...         init_method=f"tcp://:10003",
...     )
...     # Creates the destination tensordict on server.
...     # The first dim must be equal to world_size-1
...     td = TensorDict(
...         {("a", "b"): torch.zeros(2),
...          "c": torch.zeros(2)}, [2]
...     ).expand(1, 2).contiguous()
...     td.gather_and_stack(0)
...     assert td["a", "b"] != 0
...     print("yuppie")
...
>>> if __name__ == "__main__":
...     mp.set_start_method("spawn")
...
...     main_worker = mp.Process(target=server)
...     secondary_worker = mp.Process(target=client)
...
...     main_worker.start()
...     secondary_worker.start()
...
...     main_worker.join()
...     secondary_worker.join()
get(key: NestedKey, *args, as_list: bool = False, as_padded_tensor: bool = False, as_nested_tensor: bool = False, padding_side: str = 'right', layout: Optional[layout] = None, padding_value: float | int | bool = 0.0, **kwargs) Union[Tensor, TensorCollection]

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

引數:
  • key (str, str 的元組) – 要查詢的鍵。如果是 str 的元組,則等同於鏈式呼叫 getattr。

  • default

    如果 tensordict 中找不到該鍵,則返回預設值。預設為 None

    警告

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

關鍵字引數:
  • as_list (bool, optional) – 如果設定為 True,則可變長度張量將作為列表返回。與 as_padded_tensoras_nested_tensor 互斥。預設為 False

  • as_padded_tensor (bool, optional) – 如果設定為 True,則可變長度張量將作為填充張量返回。填充值可以透過 padding_value 關鍵字引數控制,填充側可以透過 padding_side 引數控制。與 as_listas_nested_tensor 互斥。預設為 False

  • as_nested_tensor (bool, optional) – 如果設定為 True,則可變長度張量將作為列表返回。與 as_listas_padded_tensor 互斥。佈局可以透過 torch.layout 引數控制。預設為 False

  • layout (torch.layout, optional) – 當 as_nested_tensor=True 時的佈局。

  • padding_side (str) – 填充側。必須是 “left”“right”。預設為 “right”

  • padding_value (scalar or bool, optional) – 填充值。預設為 0.0。

示例

>>> from tensordict import TensorDict, lazy_stack
>>> import torch
>>> td = lazy_stack([
...     TensorDict({"x": torch.ones(1,)}),
...     TensorDict({"x": torch.ones(2,) * 2}),
... ])
>>> td.get("x", as_nested_tensor=True)
NestedTensor(size=(2, j1), offsets=tensor([0, 1, 3]), contiguous=True)
>>> td.get("x", as_padded_tensor=True)
tensor([[1., 0.],
        [2., 2.]])
get_at(key: NestedKey, *args, **kwargs) Union[Tensor, TensorCollection]

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

引數:
  • key (str, tuple of str) – 要檢索的鍵。

  • index (int, slice, torch.Tensor, 可迭代物件) – 張量的索引。

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

返回:

索引的張量。

示例

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

獲取惰性堆疊中項的形狀。

異構維度將返回 -1。

此實現效率低下,因為它會嘗試堆疊各項以計算它們的形狀,並且應僅用於列印。

get_nestedtensor(key: NestedKey, default: Any = _NoDefault.ZERO, *, layout: Optional[layout] = None) Union[Tensor, TensorCollection]

當無法實現堆疊時,返回一個巢狀張量。

引數:
  • key (NestedKey) – 要巢狀的條目。

  • default (Any, optiona) –

    如果鍵不在所有子張量字典中,則返回的預設值。

    注意

    如果預設值為張量,則此方法將嘗試使用它來構建巢狀張量。否則,將返回預設值。

關鍵字引數:

layout (torch.layout, optional) – 巢狀張量的佈局。

示例

>>> td0 = TensorDict({"a": torch.zeros(4), "b": torch.zeros(4)}, [])
>>> td1 = TensorDict({"a": torch.ones(5)}, [])
>>> td = torch.stack([td0, td1], 0)
>>> a = td.get_nestedtensor("a")
>>> # using a tensor as default uses this default to build the nested tensor
>>> b = td.get_nestedtensor("b", default=torch.ones(4))
>>> assert (a == b).all()
>>> # using anything else as default returns the default
>>> b2 = td.get_nestedtensor("b", None)
>>> assert b2 is None
get_non_tensor(key: NestedKey, default=_NoDefault.ZERO)

獲取非張量值(如果存在),或者在找不到非張量值時返回 default

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

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

引數:
  • key (NestedKey) – 非張量資料的儲存位置。

  • default (Any, optional) – 找不到鍵時要返回的值。

返回: `tensordict.tensorclass.NonTensorData` 的內容,

或對應於 `key` 的條目(如果它不是 `tensordict.tensorclass.NonTensorData`)或 `default`(如果找不到條目)。

示例

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

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

half() Any

將所有張量轉換為 torch.half

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

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

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

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

引數:
  • dst (int) – 目標程序的 rank。

  • group ("ProcessGroup", optional) – 要使用的程序組。預設為 None。

  • device (torch.device, 可選) – 用於張量運算的裝置。預設為 None。

另請參閱

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

示例

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

將 TensorDict 插入到指定索引處的堆疊中。

類似於 list.insert。插入的 TensorDict 必須具有相容的 batch_size 和 device。插入是就地進行的,不返回任何內容。

引數:
  • index (int) – 新 TensorDict 應插入的索引。

  • tensordict (TensorDictBase) – 要插入到堆疊中的 TensorDict。

int() Any

將所有張量轉換為 torch.int

int16() Any

將所有張量轉換為 torch.int16

int32() Any

將所有張量轉換為 torch.int32

int64() Any

將所有張量轉換為 torch.int64

int8() Any

將所有張量轉換為 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, 可選) – 如果設定,將使用指定的程序組進行通訊。否則,將使用預設程序組。預設為 `None`。

  • return_premature (bool) – 如果為 True,則返回一個期貨列表,等待直到 tensordict 被更新。預設為 False,即在呼叫期間等待直到更新完成。

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

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

返回:

如果 return_premature=True,則返回一個未來列表以等待直到 tensordict 更新。

upon until the tensordict is updated.

is_consolidated()

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

is_contiguous() bool

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

is_empty() bool

檢查 tensordict 是否包含任何葉子。

is_floating_point() 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, return_early: bool = False) int | List['Work']

非同步傳送 tensordict 的內容。

引數:

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

關鍵字引數:
  • group (torch.distributed.ProcessGroup, 可選) – 如果設定,將使用指定的程序組進行通訊。否則,將使用預設程序組。預設為 `None`。

  • init_tag (int) – 用於標記張量的初始標籤。請注意,此標籤將根據 TensorDict 中包含的張量數量進行遞增。

  • pseudo_rand (bool) – 如果為 True,則標籤序列將是偽隨機的,允許來自不同節點的多條資料傳送而不會重疊。請注意,這些偽隨機數的生成成本很高(每個數字 1e-5 秒),這意味著它可能會減慢演算法的執行速度。預設為 False

  • return_early (bool, optional) – 如果為 True,則返回一個期貨列表,而不是最後一個傳送的張量的標籤。預設為 False

示例

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

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

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

isnan() Any

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

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

isneginf() Any

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

isposinf() Any

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

isreal() Any

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

items(include_nested=False, leaves_only=False, is_leaf=None, *, sort: bool = False)

返回 tensordict 的鍵值對生成器。

引數:
  • include_nested (bool, 可選) – 如果為 `True`,則返回巢狀值。預設為 `False`。

  • leaves_only (bool, 可選) – 如果為 `False`,則僅返回葉子節點。預設為 `False`。

  • is_leaf (callable, optional) –

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

    注意

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

    另請參閱

    is_leaf_nontensor()default_is_leaf()

關鍵字引數:

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

keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) _LazyStackedTensorDictKeysView

返回 tensordict 鍵的生成器。

警告

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

引數:
  • include_nested (bool, 可選) – 如果為 `True`,則返回巢狀值。預設為 `False`。

  • leaves_only (bool, 可選) – 如果為 `False`,則僅返回葉子節點。預設為 `False`。

  • is_leaf (callable, optional) –

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

    注意

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

    另請參閱

    is_leaf_nontensor()default_is_leaf()

關鍵字引數:

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

示例

>>> from tensordict import TensorDict
>>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[])
>>> data.keys()
['0', '1']
>>> list(data.keys(leaves_only=True))
['0']
>>> list(data.keys(include_nested=True, leaves_only=True))
['0', '1', ('1', '2')]
classmethod lazy_stack(items: Sequence[TensorCollection], dim: int = 0, *, device: Optional[Union[device, str, int]] = None, out: Optional[T] = None, stack_dim_name: Optional[str] = None, **kwargs) T

將 tensordicts 堆疊在 LazyStackedTensorDict 中。

引數:
  • items (Sequence of TensorDictBase instances) – 要堆疊的 TensorDictBase 例項序列。

  • dim (int, optional) – 進行惰性堆疊的維度。預設為 0。

關鍵字引數:
  • device (torch.device, optional) – 在 LazyStackedTensorDict 中設定的裝置,以防無法從 tensordict 列表中推斷出來(例如,列表為空)。

  • out (TensorDictBase, optional) – 要將資料寫入的 LazyStackedTensorDict

  • stack_dim_name (str, optional) – 堆疊維度的名稱。

  • strict_shape (bool, optional) – 如果為 True,則必須匹配每個 tensordict 的形狀。預設為 False

  • as_list (bool, optional) – 如果設定為 True,則可變長度張量將作為列表返回。與 as_padded_tensoras_nested_tensor 互斥。預設為 False

  • as_padded_tensor (bool, optional) – 如果設定為 True,則可變長度張量將作為填充張量返回。填充值可以透過 padding_value 關鍵字引數控制,填充側可以透過 padding_side 引數控制。與 as_listas_nested_tensor 互斥。預設為 False

  • as_nested_tensor (bool, optional) – 如果設定為 True,則可變長度張量將作為列表返回。與 as_listas_padded_tensor 互斥。佈局可以透過 torch.layout 引數控制。預設為 False

  • layout (torch.layout, optional) – 當 as_nested_tensor=True 時的佈局。

  • padding_side (str) – 填充側。必須是 “left”“right”。預設為 “right”

  • padding_value (scalar or bool, optional) – 填充值。預設為 0.0。

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

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

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

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

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

原地版本的 lerp()

lgamma() Any

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

lgamma_() Any

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

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

從磁碟載入 tensordict。

此類方法是 load_memmap() 的代理。

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

在當前 tensordict 中從磁碟載入 tensordict。

此類方法是 load_memmap_() 的代理。

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

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

引數:
  • prefix (str資料夾路徑) – 應從中獲取已儲存 tensordict 的資料夾路徑。

  • device (torch.device等效項, 可選) – 如果提供,資料將非同步轉換為該裝置。支援 `"meta"` 裝置,在這種情況下,資料不會被載入,而是建立一組空的 "meta" 張量。這對於在不實際開啟任何檔案的情況下了解模型大小和結構很有用。

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

  • out (TensorDictBase, 可選) – 應將資料寫入其中的可選 tensordict。

示例

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

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

示例

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

tensordict 也可以在“meta”裝置上載入,或者作為假張量載入。

示例

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

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

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

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

載入一個 state-dict(格式如 state_dict())到 tensordict 中。

引數:
  • state_dict (OrderedDict) – 要複製的 state_dict。

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

  • assign (bool, optional) – 是否將 state 字典中的項分配給 tensordict 中相應的鍵,而不是就地複製到 tensordict 的當前張量中。當 False 時,將保留當前模組張量的屬性;當 True 時,將保留 state 字典中張量的屬性。預設為 False

  • from_flatten (bool, optional) – 如果為 True,則假定輸入的 state_dict 是扁平化的。預設為 False

示例

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

鎖定 tensordict 以進行非就地操作。

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

此方法可用作裝飾器。

示例

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

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

log10() Any

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

log10_() Any

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

log1p() Any

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

log1p_() Any

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

log2() Any

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

log2_() Any

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

log_() Any

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

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

selfother 執行邏輯 AND 操作。

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

other (TensorDictBase or torch.Tensor) – 用於執行邏輯 AND 操作的張量或 TensorDict。

關鍵字引數:

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

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

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

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

引數:
  • dim (int or tuple of ints, optional) – 要約簡的維度或維度。如果為 None,則約簡 tensordict 的所有批次維度。

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

關鍵字引數:

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

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

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

警告

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

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

引數:
  • key (NestedKey) – 要寫入的新條目的鍵。如果鍵已存在於 tensordict 中,將引發異常。

  • shape (torch.Size or equivalent, torch.Tensor for nested tensors) – 要寫入的張量的形狀。

關鍵字引數:

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

返回:

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

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

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

警告

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

注意

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

引數:
  • key (NestedKey) – 要寫入的新條目的鍵。如果鍵已存在於 tensordict 中,將引發異常。

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

  • shape (torch.Size or equivalent, torch.Tensor for nested tensors) – 要寫入的張量的形狀。

關鍵字引數:

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

返回:

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

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

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

警告

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

如果 copy_dataTrue(即儲存是共享的),此方法將始終複製儲存內容。

引數:
  • key (NestedKey) – 要寫入的新條目的鍵。如果鍵已存在於 tensordict 中,將引發異常。

  • tensor (torch.Tensor) – 要在物理記憶體中複製的張量。

關鍵字引數:

copy_data (bool, optionaL) – 如果為 False,則新張量將共享輸入(如形狀和 dtype)的元資料,但內容將為空。預設為 True

返回:

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

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

將一個函式對映到tensordict在某個維度上的切片。

此方法將透過將 tensordict 切塊成相等大小的 tensordict 並分派操作到所需的 Worker 數量來應用一個函式。

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

注意

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

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

  • dim (int, 可選) – tensordict 將被分塊的維度。

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

關鍵字引數:
  • out (TensorDictBase, 可選) – 一個可選的輸出容器。其沿提供的 dim 的批次大小必須與 self.ndim 匹配。如果它是共享的或記憶體對映的(is_shared()is_memmap() 返回 True),它將在遠端程序中被填充,避免資料內向傳輸。否則,來自 self 切片的將資料傳送到程序,在當前程序中收集,然後就地寫入 out

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

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

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

  • generator (torch.Generator, optional) –

    一個用於播種的生成器。將從此生成器生成一個基礎種子,並且池中的每個工作程序將使用提供的種子(並加上從 0num_workers 的唯一整數)進行播種。如果未提供生成器,將使用一個隨機整數作為種子。要與未播種的工作程序一起工作,應單獨建立池並將其傳遞給 map()

    注意

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

    注意

    為工作執行緒設定種子的目的是讓每個工作執行緒擁有不同的種子,而不是為了讓 map 方法的呼叫結果可重現。換句話說,兩個實驗可能並且很可能返回不同的結果,因為無法知道哪個工作執行緒會執行哪個作業。但是,我們可以確保每個工作執行緒都有不同的種子,並且每個工作執行緒上的偽隨機操作將是不相關的。

  • max_tasks_per_child (int, 可選) – 每個子程序拾取的最大任務數。預設為 None,即對任務數量沒有限制。

  • worker_threads (int, 可選) – 工作程序的執行緒數。預設為 1

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

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

  • mp_start_method (str, optional) – 多程序的啟動方法。如果未提供,將使用預設啟動方法。可接受的字串是 "fork""spawn"。請注意,使用 "fork" 啟動方法時,"cuda" 張量無法在程序之間共享。如果 pool 引數傳遞給 map 方法,則此引數無效。

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> def process_data(data):
...     data.set("y", data.get("x") + 1)
...     return data
>>> if __name__ == "__main__":
...     data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_()
...     data = data.map(process_data, dim=1)
...     print(data["y"][:, :10])
...
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
map_iter(fn: Callable[[TensorCollection], TensorCollection | None], dim: int = 0, num_workers: int | None = None, *, shuffle: bool = False, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: torch.Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = True, pbar: bool = False, mp_start_method: str | None = None) Iterator[T]

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

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

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

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

注意

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

注意

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

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

  • dim (int, 可選) – tensordict 將被分塊的維度。

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

關鍵字引數:
  • shuffle (bool, 可選) – 索引是否應全域性混洗。如果為 True,則每個批次將包含非連續的樣本。如果 index_with_generator=Falseshuffle=True`,則會引發錯誤。預設為 False

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

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

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

  • generator (torch.Generator, optional) –

    一個用於播種的生成器。將從此生成器生成一個基礎種子,並且池中的每個工作程序將使用提供的種子(並加上從 0num_workers 的唯一整數)進行播種。如果未提供生成器,將使用一個隨機整數作為種子。要與未播種的工作程序一起工作,應單獨建立池並將其傳遞給 map()

    注意

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

    注意

    為工作執行緒設定種子的目的是讓每個工作執行緒擁有不同的種子,而不是為了讓 map 方法的呼叫結果可重現。換句話說,兩個實驗可能並且很可能返回不同的結果,因為無法知道哪個工作執行緒會執行哪個作業。但是,我們可以確保每個工作執行緒都有不同的種子,並且每個工作執行緒上的偽隨機操作將是不相關的。

  • max_tasks_per_child (int, 可選) – 每個子程序拾取的最大任務數。預設為 None,即對任務數量沒有限制。

  • worker_threads (int, 可選) – 工作程序的執行緒數。預設為 1

  • index_with_generator (bool, 可選) –

    如果為 True,則將在查詢時進行 tensordict 的拆分/分塊,從而節省初始化時間。請注意,chunk()split() 比索引(在生成器中使用)效率更高,因此初始化時間的處理時間收益可能會對總執行時間產生負面影響。預設為 True

    注意

    index_with_generator 的預設值對於 map_itermap 是不同的,前者假定將拆分後的 TensorDict 儲存在記憶體中成本過高。

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

  • mp_start_method (str, optional) – 多程序的啟動方法。如果未提供,將使用預設啟動方法。可接受的字串是 "fork""spawn"。請注意,使用 "fork" 啟動方法時,"cuda" 張量無法在程序之間共享。如果 pool 引數傳遞給 map 方法,則此引數無效。

示例

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

masked_fill 的非原地版本。

引數:
  • mask (boolean torch.Tensor) – 要填充的值的掩碼。形狀必須與 tensordict 的批次大小匹配。

  • value – 用於填充張量的值。

返回:

self

示例

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

用期望值填充與掩碼對應的項。

引數:
  • mask (boolean torch.Tensor) – 要填充的值的掩碼。形狀必須與 tensordict 的批次大小匹配。

  • value – 用於填充張量的值。

返回:

self

示例

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

遮蔽 TensorDict 的所有張量,並返回一個具有指向被遮蔽值的新 TensorDict 例項。

引數:

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

示例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)},
...    batch_size=[3])
>>> mask = torch.tensor([True, False, False])
>>> td_mask = td.masked_select(mask)
>>> td_mask.get("a")
tensor([[0., 0., 0., 0.]])
max(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None, return_indices: bool = True) Self | torch.Tensor

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

引數:
  • dim (int, optional) – 如果為 None,則返回一個無維度的 tensordict,其中包含所有葉子的最大值(如果可以計算)。如果為整數,則僅當該維度與 tensordict 的形狀相容時,才會在指定的維度上呼叫 max

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

關鍵字引數:
  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

  • return_argmins (bool, 可選) – max() 在傳遞 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.max(dim=0)
max(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.max()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.max(reduce=True)
tensor(3.2942)
maximum(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) Any

計算 selfother 的逐元素最大值。

引數:

other (TensorDictTensor) – 另一個輸入tensordict或張量。

關鍵字引數:

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

maximum_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) Any

maximum() 的就地版本。

注意

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

classmethod maybe_dense_stack(items: Sequence[TensorDictBase], dim: int = 0, out: Optional[T] = None, strict: bool = False) Any

如果可能,密集地堆疊張量或 tensordict,否則堆疊到 LazyStackedTensorDict。

示例

>>> td0 = TensorDict({"a": 0}, [])
>>> td1 = TensorDict({"b": 0}, [])
>>> LazyStackedTensorDict.maybe_dense_stack([td0, td0])  # returns a TensorDict with shape [2]
>>> LazyStackedTensorDict.maybe_dense_stack([td0, td1])  # returns a LazyStackedTensorDict with shape [2]
>>> LazyStackedTensorDict.maybe_dense_stack(list(torch.randn(2)))  # returns a torch.Tensor with shape [2]
maybe_to_stack()

非張量堆疊與非堆疊之間的可互換性佔位符。

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

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

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

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

關鍵字引數:
  • dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定,則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設:None

  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.mean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.mean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.mean(reduce=True)
tensor(-0.0547)
>>> td.mean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.mean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.mean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])
memmap(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) Any

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

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

  • copy_existing (bool) – 如果為 False(預設值),並且 tensordict 中某項已是儲存在磁碟上的張量且關聯了檔案,但未按 prefix 儲存到正確位置,則會引發異常。如果為 True,則任何現有張量都將被複制到新位置。

關鍵字引數:
  • num_threads (int, 可選) – 用於寫入 memmap 張量的執行緒數。預設為 0

  • return_early (bool, 可選) – 如果為 Truenum_threads>0,則方法將返回一個 tensordict 的 future。

  • share_non_tensor (bool, 可選) – 如果為 True,則非張量資料將在程序之間共享,並且在單個節點內的任何工作程序上進行寫入操作(例如就地更新或設定)將更新所有其他工作程序上的值。如果非張量葉子的數量很高(例如,共享大型非張量資料堆疊),這可能會導致 OOM 或類似錯誤。預設為 False

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

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

返回:

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

注意

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

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

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

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

  • copy_existing (bool) – 如果為 False(預設值),並且 tensordict 中某項已是儲存在磁碟上的張量且關聯了檔案,但未按 prefix 儲存到正確位置,則會引發異常。如果為 True,則任何現有張量都將被複制到新位置。

關鍵字引數:
  • num_threads (int, 可選) – 用於寫入 memmap 張量的執行緒數。預設為 0

  • return_early (bool, optional) – 如果為 Truenum_threads>0,則該方法將返回一個 tensordict 的 future。可以透過 future.result() 查詢結果 tensordict。

  • share_non_tensor (bool, 可選) – 如果為 True,則非張量資料將在程序之間共享,並且在單個節點內的任何工作程序上進行寫入操作(例如就地更新或設定)將更新所有其他工作程序上的值。如果非張量葉子的數量很高(例如,共享大型非張量資料堆疊),這可能會導致 OOM 或類似錯誤。預設為 False

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

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

返回:

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

注意

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

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

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

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

  • copy_existing (bool) – 如果為 False(預設值),並且 tensordict 中某項已是儲存在磁碟上的張量且關聯了檔案,但未按 prefix 儲存到正確位置,則會引發異常。如果為 True,則任何現有張量都將被複制到新位置。

關鍵字引數:
  • num_threads (int, 可選) – 用於寫入 memmap 張量的執行緒數。預設為 0

  • return_early (bool, 可選) – 如果為 Truenum_threads>0,則方法將返回一個 tensordict 的 future。

  • share_non_tensor (bool, 可選) – 如果為 True,則非張量資料將在程序之間共享,並且在單個節點內的任何工作程序上進行寫入操作(例如就地更新或設定)將更新所有其他工作程序上的值。如果非張量葉子的數量很高(例如,共享大型非張量資料堆疊),這可能會導致 OOM 或類似錯誤。預設為 False

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

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

返回:

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

注意

這是將一組大型緩衝區寫入磁碟的推薦方法,因為 memmap_() 將複製資訊,這對於大型內容來說可能很慢。

示例

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

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

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

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

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

引數:
  • dim (int, optional) – 如果 None,則返回一個無維度的 tensordict,其中包含所有葉子的最小值(如果可以計算)。如果為整數,則僅當此維度與 tensordict 的形狀相容時,才會對指定的維度呼叫 min

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

關鍵字引數:
  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

  • return_argmins (bool, 可選) – 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._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) Any

計算 selfother 的逐元素最小值。

引數:

other (TensorDictTensor) – 另一個輸入tensordict或張量。

關鍵字引數:

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

minimum_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) Any

minimum() 的就地版本。

注意

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

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

other 乘以 self

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

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

引數:

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

關鍵字引數:

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

mul_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) Any

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[Any]

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

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

引數:
  • fn (Callable) – 要應用於 tensordict 中的(名稱,張量)對的函式。對於每個葉子,只使用其葉子名稱(不使用完整的 NestedKey)。

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

  • nested_keys (bool, optional) – 如果為 True,則將使用葉子的完整路徑。預設為 False,即僅將最後一個字串傳遞給函式。

  • batch_size (sequence of int, optional) – 如果提供,結果 TensorDict 將具有所需的 batch_size。batch_size 引數應與轉換後的 batch_size 匹配。這是一個僅限關鍵字的引數。

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

  • names (字串列表, 可選) – 新的維度名稱,以防batch_size被修改。

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

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

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

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

  • call_on_nested (bool, optional) –

    如果為 True,則該函式將應用於第一級張量和容器(TensorDict 或 tensorclass)。在這種情況下,func 負責傳播其呼叫到巢狀級別。這允許在將呼叫傳播到巢狀 tensordicts 時進行精細控制。如果為 False,則該函式僅應用於葉子節點,並且 apply 將負責將函式分派到所有葉子節點。

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

  • out (TensorDictBase, 可選) –

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

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

    警告

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

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

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

返回:

一個包含轉換後張量的新tensordict。

示例

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

注意

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

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

tensordict 的維度名稱。

可以使用 names 引數在構造時設定名稱。

另請參閱 refine_names() 以瞭解如何構造後設置名稱的詳細資訊。

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

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

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

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

關鍵字引數:
  • dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定,則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設:None

  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.nanmean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.nanmean(reduce=True)
tensor(-0.0547)
>>> td.nanmean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.nanmean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])
nansum(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) Union[Any, Tensor]

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

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

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

關鍵字引數:
  • dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定,則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設:None

  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

示例

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

請參閱 batch_dims()

ndimension() int

請參閱 batch_dims()

neg() Any

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

neg_() Any

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

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

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

預設情況下,返回的 TensorDict 具有與此 tensordict 相同的 torch.dtypetorch.device

引數:

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

關鍵字引數:
  • dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值:如果 None,則 torch.dtype 將保持不變。

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

  • requires_grad (bool, optional) – 是否應該為返回的張量記錄autograd操作。預設值:False

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

  • pin_memory (bool, optional) – 如果設定為True,返回的張量將被分配到固定記憶體中。僅適用於CPU張量。預設值:False

  • empty_lazy (bool, optional) – 如果為True,則懶惰堆疊的內容將被清空。當新tensordict在填充過程中懶惰堆疊的內容可能發生變化時,這可能很有用。此引數將傳播到子tensordict。預設為False

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

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

預設情況下,返回的 TensorDict 具有與此 tensordict 相同的 torch.dtypetorch.device

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

  • fill_value (scalar) – 用於填充輸出 Tensor 的數值。

關鍵字引數:
  • dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值:如果 None,則 torch.dtype 將保持不變。

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

  • requires_grad (bool, optional) – 是否應該為返回的張量記錄autograd操作。預設值:False

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

  • pin_memory (bool, optional) – 如果設定為True,返回的張量將被分配到固定記憶體中。僅適用於CPU張量。預設值:False

  • empty_lazy (bool, optional) – 如果為True,則懶惰堆疊的內容將被清空。當新tensordict在填充過程中懶惰堆疊的內容可能發生變化時,這可能很有用。此引數將傳播到子tensordict。預設為False

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

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

預設情況下,返回的 TensorDict 具有與此 tensordict 相同的 torch.dtypetorch.device

引數:

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

關鍵字引數:
  • dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值:如果 None,則 torch.dtype 將保持不變。

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

  • requires_grad (bool, optional) – 是否應該為返回的張量記錄autograd操作。預設值:False

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

  • pin_memory (bool, optional) – 如果設定為True,返回的張量將被分配到固定記憶體中。僅適用於CPU張量。預設值:False

  • empty_lazy (bool, optional) – 如果為True,則懶惰堆疊的內容將被清空。當新tensordict在填充過程中懶惰堆疊的內容可能發生變化時,這可能很有用。此引數將傳播到子tensordict。預設為False

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

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

預設情況下,返回的 TensorDict 值具有與此張量相同的 torch.dtypetorch.device

data (torch.TensorTensorDictBase) – 要複製的資料。

引數:

data (torch.TensorTensorDictBase) – 要複製的資料。

關鍵字引數:
  • dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值:如果 None,則 torch.dtype 將保持不變。

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

  • requires_grad (bool, optional) – 是否應該為返回的張量記錄autograd操作。預設值:False

  • pin_memory (bool, optional) – 如果設定為True,返回的張量將被分配到固定記憶體中。僅適用於CPU張量。預設值:False

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

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

預設情況下,返回的 TensorDict 具有與此 tensordict 相同的 torch.dtypetorch.device

引數:

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

關鍵字引數:
  • dtype (torch.dtype, optional) – 所需返回 tensordict 的型別。預設值:如果 None,則 torch.dtype 將保持不變。

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

  • requires_grad (bool, optional) – 是否應該為返回的張量記錄autograd操作。預設值:False

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

  • pin_memory (bool, optional) – 如果設定為True,返回的張量將被分配到固定記憶體中。僅適用於CPU張量。預設值:False

  • empty_lazy (bool, optional) – 如果為True,則懶惰堆疊的內容將被清空。當新tensordict在填充過程中懶惰堆疊的內容可能發生變化時,這可能很有用。此引數將傳播到子tensordict。預設為False

non_tensor_items(include_nested: bool = False)

遞迴地返回所有非張量葉節點。

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

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

關鍵字引數:
  • out (TensorDict, optional) – 輸出 tensordict。

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

numel() int

批次中的總元素數量。

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

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

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

非張量資料按原樣公開。

示例

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

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

param_count(*, count_duplicates: bool = True) int

計算引數數量(可索引項的總數),僅考慮張量。

關鍵字引數:

count_duplicates (bool) – 是否計算重複的張量作為獨立的張量。如果為 False,則只丟棄完全相同的張量(來自同一基礎張量的相同檢視但具有不同 ID 的張量將被計算兩次)。預設為 True(每個張量假定為單個副本)。

permute(*args, **kwargs)

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

引數:
  • *dims_list (int) – tensordict 的批處理維度的順序。或者,可以提供單個整數可迭代物件。

  • dims (list of int) – 呼叫 permute(…) 的替代方法。

返回:

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

示例

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

在儲存的張量上呼叫 pin_memory()

引數:
  • num_threads (intstr) – 如果提供,則呼叫pin_memory時使用的執行緒數。預設為None,它在ThreadPoolExecutor(max_workers=None)中設定了大量執行緒。要在主執行緒上執行所有呼叫pin_memory(),請傳入num_threads=0

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

pin_memory_(num_threads: int | str = 0) Any

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

引數:

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

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

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

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

引數:
  • key (str or nested key) – 要查詢的條目。

  • default (Any, optional) – 找不到鍵時要返回的值。

示例

>>> td = TensorDict({"1": 1}, [])
>>> one = td.pop("1")
>>> assert one == 1
>>> none = td.pop("1", default=None)
>>> assert none is None
popitem() Tuple[NestedKey, Union[Tensor, TensorCollection]]

移除最後插入 TensorDict 的項。

popitem 只會返回非巢狀值。

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

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

other 可以是單個 float 數、TensorTensorDict

other 是張量時,inputother 的形狀必須是可廣播的。

引數:

other (float, tensor or tensordict) – 指數

關鍵字引數:

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

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

原地版本 pow()

注意

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

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

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

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

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

關鍵字引數:
  • dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定,則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設:None

  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

示例

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

將所有張量強制轉換為 torch.qint32

qint8() Any

將所有張量強制轉換為 torch.qint8

quint4x2() Any

將所有張量強制轉換為 torch.quint4x2

quint8() Any

將所有張量強制轉換為 torch.quint8

reciprocal() Any

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

reciprocal_() Any

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

record_stream(stream: Stream) Any

標記 TensorDict 已被此流使用。

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

更多資訊請參閱 record_stream()

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

接收 TensorDict 的內容並用其更新內容。

有關上下文,請檢視 send 方法中的示例。

引數:

src (int) – 源工作程序的 rank。

關鍵字引數:
  • group (torch.distributed.ProcessGroup, 可選) – 如果設定,將使用指定的程序組進行通訊。否則,將使用預設程序組。預設為 `None`。

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

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

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

在所有機器上歸約 TensorDict。

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

refine_names(*names) Any

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

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

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

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

Returns: 具有根據輸入命名的維度的相同的 TensorDict。

示例

>>> td = TensorDict({}, batch_size=[3, 4, 5, 6])
>>> tdr = td.refine_names(None, None, None, "d")
>>> assert tdr.names == [None, None, None, "d"]
>>> tdr = td.refine_names("a", None, None, "d")
>>> assert tdr.names == ["a", None, None, "d"]
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) Any

用新字串重新命名一個鍵,並返回具有更新後的鍵名的相同 TensorDict。

引數:
  • old_key (strnested key) – 要重新命名的鍵。

  • new_key (strnested key) – 條目的新名稱。

  • safe (bool, optional) – 如果為 True,則當新鍵已存在於 TensorDict 中時將引發錯誤。

返回:

self

repeat(*repeats: int) Any

沿指定維度重複此張量。

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

警告

repeat() 的行為與 numpy.repeat() 不同,但更接近 numpy.tile()。要獲得與 numpy.repeat() 類似的運算子,請參閱 repeat_interleave()

引數:

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

示例

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

重複 TensorDict 的元素。

警告

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

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

  • dim (int, optional) – 重複值的維度。預設為使用展平的輸入陣列,並返回一個展平的輸出陣列。

關鍵字引數:

output_size (int, optional) – 給定軸的總輸出大小(例如,重複的總和)。如果給定,它將避免需要同步流來計算tensordict的輸出形狀。

返回:

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

示例

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

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

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

返回:

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

requires_grad_(requires_grad=True) Any

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

返回此 TensorDict。

引數:

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

reshape(*args, **kwargs) Any

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

引數:

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

返回:

具有重塑鍵的 TensorDict。

示例

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

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

round_() Any

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

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

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

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

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

引數:

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

關鍵字引數:
  • alpha (Number) – other 的乘數。

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

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

將tensordict儲存到磁碟。

此函式是 memmap() 的代理。

property saved_path

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

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

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

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

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

引數:
  • *keys (str) – 要選擇的鍵。

  • inplace (bool) – 如果為 True,則 tensordict 將就地剪枝。預設為 False

  • strict (bool, optional) – 選擇一個不存在的鍵是否會返回錯誤。預設為 True

返回:

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

注意

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

示例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, [])
>>> td.select("a", ("b", "c"))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.select("a", "b")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.select("this key does not exist", strict=False)
TensorDict(
    fields={
    },
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
send(dst: int, *, group: 'torch.distributed.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) None

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

引數:

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

關鍵字引數:
  • group (torch.distributed.ProcessGroup, 可選) – 如果設定,將使用指定的程序組進行通訊。否則,將使用預設程序組。預設為 `None`。

  • init_tag (int) – 用於標記張量的初始標籤。請注意,此標籤將根據 TensorDict 中包含的張量數量進行遞增。

  • pseudo_rand (bool) – 如果為 True,則標籤序列將是偽隨機的,允許來自不同節點的多條資料傳送而不會重疊。請注意,這些偽隨機數的生成成本很高(每個數字 1e-5 秒),這意味著它可能會減慢演算法的執行速度。預設為 False

示例

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

將指定的鍵從 tensordict 中分離(就地操作)。

另請參閱

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

另請參閱

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

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

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

  • strict (bool, optional) – 如果為 True,則在鍵丟失時引發異常。預設為 True

  • filter_empty (bool, optional) – 如果為 True,則會刪除 self 中的空 tensordicts。預設為 True

返回:

分離出的 tensordict。

返回型別:

T

示例

>>> td = TensorDict(
...     a=0,
...     b=0,
...     c=0,
...     d=0,
... )
>>> td_a_c = td.separates("a", "c")
>>> print(td_a_c)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(td)
TensorDict(
    fields={
        b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
set(key: NestedKey, item: Union[Tensor, TensorCollection], inplace: bool = False, *, non_blocking: bool = False, **kwargs: Any) Any

設定一個新的鍵值對。

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

  • item (torch.Tensor or equivalent, TensorDictBase instance) – 要儲存在 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: Union[Tensor, TensorCollection], *, non_blocking: bool = False) Any

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

引數:
關鍵字引數:

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: Union[Tensor, TensorCollection], index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], *, non_blocking: bool = False) Any

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

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

  • value (torch.Tensor) – 在 index 處設定的值。

  • index (int, tensor or 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.tensorclass.NonTensorData 在 tensordict 中註冊一個非張量值。

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

返回:self

示例

>>> data = TensorDict({}, batch_size=[])
>>> data.set_non_tensor(("nested", "the string"), "a string!")
>>> assert data.get_non_tensor(("nested", "the string")) == "a string!"
>>> # regular `get` works but returns a NonTensorData object
>>> data.get(("nested", "the string"))
NonTensorData(
    data='a string!',
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
setdefault(key: NestedKey, default: Union[Tensor, TensorCollection, Any], inplace: bool = False) Union[Tensor, TensorCollection]

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

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

引數:
  • key (strnested key) – 值的名稱。

  • default (torch.Tensor or compatible type, TensorDictBase) – 如果鍵不存在,則儲存在 tensordict 中的值。

返回:

tensordict 中 key 的值。如果 key 之前未設定,則為 default。

示例

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

請參閱 batch_size

share_memory_() Any

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

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

返回:

self

sigmoid() Any

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

sigmoid_() Any

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

sign() Any

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

sign_() Any

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

sin() Any

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

sin_() Any

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

sinh() Any

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

sinh_() Any

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

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

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

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

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

將 softmax 函式應用於 tensordict 元素。

引數:
  • dim (intint 元組) – 將在其上計算 softmax 的 tensordict 維度。

  • dtype (torch.dtype, optional) – 返回張量的期望資料型別。如果指定,則在執行操作之前,輸入張量將被轉換為 dtype。這有助於防止資料型別溢位。

property sorted_keys: list[tensordict._nestedkey.NestedKey]

按字母順序返回鍵。

不支援其他引數。

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

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

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

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

引數:
  • split_size (int or List(int)) – 單個塊的大小或每個塊的大小列表。

  • dim (int) – 沿此維度分割張量。

返回:

一個具有給定維度中指定尺寸的 TensorDict 列表。

示例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td0, td1 = td.split([1, 2], dim=0)
>>> print(td0['x'])
torch.Tensor([[0, 1, 2, 3]])
split_keys(*key_sets, inplace=False, default: Any = _NoDefault.ZERO, strict: bool = True, reproduce_struct: bool = False) Tuple[T, ...]

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

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

引數:
  • key_sets (鍵的序列 Dict[in_key, out_key] 或 鍵列表) – 各個分割。

  • inplace (bool, optional) – 如果為 True,則 self 中的鍵將被原地移除。預設為 False

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

  • strict (bool, optional) – 如果為 True,則在鍵丟失時引發異常。預設為 True

  • reproduce_struct (bool, optional) – 如果為 True,則所有返回的 tensordict 都將具有與 self 相同的樹結構,即使某些子 tensordict 不包含葉子。

注意

None 非張量值將被忽略且不返回。

注意

該方法不檢查提供的列表中的重複項。

示例

>>> td = TensorDict(
...     a=0,
...     b=0,
...     c=0,
...     d=0,
... )
>>> td_a, td_bc, td_d = td.split_keys(["a"], ["b", "c"])
>>> print(td_bc)
sqrt() Any

計算 self 的逐元素平方根。

sqrt_() Any

平方根的原地版本。

squeeze(*args, **kwargs)

壓縮維度在 -self.batch_dims+1self.batch_dims-1 之間的所有張量,並將它們返回到一個新的 tensordict 中。

引數:

dim (Optional[int]) – 壓縮的維度。如果 dim 為 None,則會壓縮所有單例維度。預設為 None

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> td = td.squeeze()
>>> td.shape
torch.Size([3, 4])
>>> td.get("x").shape
torch.Size([3, 4, 2])

此操作也可以作為上下文管理器使用。對原始 tensordict 的更改將是異變的,即原始張量的內容不會被修改。這也假定 tensordict 未被鎖定(否則,需要解鎖 tensordict)。此功能與隱式壓縮不相容。

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> with td.squeeze(1) as tds:
...     tds.set("y", torch.zeros(3, 4))
>>> assert td.get("y").shape == [3, 1, 4]
classmethod stack(input, dim: int = 0, *, out=None)

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

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

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

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

引數:

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

關鍵字引數:
  • sorted (boolNestedKeys 列表) – 如果 True,條目將按字母順序堆疊。如果 False(預設),則使用字典順序。或者,可以提供鍵名列表,並按此順序堆疊張量。這會產生一些開銷,因為鍵列表將與tensordict中的葉節點名稱列表進行比較。

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

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

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

引數:

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

關鍵字引數:
  • out_key (NestedKey) – 堆疊輸入的新的鍵名。

  • keep_entries (bool, optional) – 如果為 False,則 keys 中的條目將被刪除。預設為 False

  • dim (int, 可選) – 必須進行堆疊的維度。預設為 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。

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

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

引數:
  • destination (dict, optional) – 如果提供,tensordict 的狀態將更新到 dict 中,並返回相同的物件。否則,將建立一個 OrderedDict 並返回。預設:None

  • prefix (str, optional) – 新增到張量名稱的字首,用於組合 state_dict 中的鍵。預設:''

  • keep_vars (bool, optional) – 預設情況下,state_dict 中返回的 torch.Tensor 項將從 autograd 中分離。如果設定為 True,則不會執行分離。預設:False

  • flatten (bool, optional) – 是否應使用 "." 字元展平結構。預設為 False

示例

>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, [])
>>> sd = data.state_dict()
>>> print(sd)
OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3', OrderedDict([('3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])), ('__batch_size', torch.Size([])), ('__device', None)])
>>> sd = data.state_dict(flatten=True)
OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3.3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])
std(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) Union[Any, Tensor]

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

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

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

關鍵字引數:
  • correction (int) – 樣本大小與樣本自由度之差。預設為貝塞爾校正,correction=1。

  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.std(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.std()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.std(reduce=True)
tensor(1.0006)
>>> td.std(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.std(reduce=True, dim="feature")
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
>>> td.std(reduce=True, dim=0)
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])
sub(other: tensordict.base.TensorDictBase | torch.Tensor | float, *, alpha: float | None = None, default: Optional[Union[str, Tensor, TensorCollection]] = None) Any

other 縮放 alpha 後,從 self 中減去。

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

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

引數:

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

關鍵字引數:
  • alpha (Number) – other 的乘數。

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

sub_(other: tensordict.base.TensorDictBase | torch.Tensor | float, alpha: float | None = None)

原地版本 sub()

注意

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

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

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

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

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

關鍵字引數:
  • dtype (torch.dtype, optional) – 返回張量所需的 dtype。如果指定,則在執行操作之前將輸入張量轉換為 dtype。這有助於防止資料型別溢位。預設:None

  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

示例

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

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

tan_() Any

原地計算 TensorDict 中每個元素的 tan() 值。

tanh() Any

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

tanh_() Any

原地計算 TensorDict 中每個元素的 tanh() 值。

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

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

引數:
  • indices_or_sections (intList(int) 或 tuple(int) 或 一維整數張量) – 如果 indices_or_sections 是一個整數 n 或一個值為 n 的零維長整型張量,則輸入沿維度 dim 分成 n 部分。如果沿維度 dim 輸入可被 n 整除,則每部分的大小均為 input.size(dim) / n。如果輸入不能被 n 整除,則前 int(input.size(dim) % n) 部分的大小為 int(input.size(dim) / n) + 1,其餘部分大小為 int(input.size(dim) / n)。如果 indices_or_sections 是整數列表或元組,或一維長整型張量,則輸入沿維度 dim 在列表、元組或張量中的每個索引處分割。例如,indices_or_sections=[2, 3]dim=0 將產生張量 input[:2]input[2:3]input[3:]。如果 indices_or_sections 是張量,則它必須是 CPU 上的零維或一維長整型張量。

  • dim (int, optional) – 要分割張量的維度。預設為 0

示例

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

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

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

引數:
  • device (torch.device, optional) – tensordict 的期望裝置。

  • dtype (torch.dtype, optional) – tensordict 的期望浮點數或複數 dtype。

  • tensor (torch.Tensor, optional) – 此 TensorDict 中所有張量的期望 dtype 和裝置都取自此張量。

關鍵字引數:
  • non_blocking (bool, 可選) – 操作是否應為阻塞的。

  • memory_format (torch.memory_format, optional) – 此 tensordict 中 4D 引數和緩衝區所需的記憶體格式。

  • batch_size (torch.Size, optional) – 輸出 tensordict 的結果 batch-size。

  • other (TensorDictBase, optional) –

    TensorDict 例項,其 dtype 和 device 是該 TensorDict 中所有張量的目標 dtype 和 device。

    注意

    由於 TensorDictBase 例項沒有 dtype,因此 dtype 是從示例葉子節點收集的。如果存在多個 dtype,則不會進行 dtype 轉換。

  • non_blocking_pin (bool, 可選) –

    如果為 True,則張量將在傳送到裝置之前被 pinned。這將非同步進行,但可以透過 num_threads 引數進行控制。

    注意

    tensordict.pin_memory().to("cuda") 通常比 tensordict.to("cuda", non_blocking_pin=True) 慢,因為在後一種情況下,pin_memory 是非同步呼叫的。多執行緒 pin_memory 通常在張量大且數量多時很有益:當張量數量太少時,建立執行緒和收集資料的開銷會超過多執行緒的優勢;如果張量很小,遍歷長列表的開銷也可能過大。

  • num_threads (intNone, 可選) – 如果 non_blocking_pin=True,則用於 pin_memory 的執行緒數。預設情況下,將生成 max(1, torch.get_num_threads()) 個執行緒。num_threads=0 將取消 pin_memory() 呼叫中的所有多執行緒。

  • inplace (bool, 可選) – 如果為 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, convert_tensors: Union[bool, Literal['numpy']] = False, tolist_first: bool = False, as_linked_list: bool = False) dict[str, Any]

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

引數:
  • retain_none (bool) – 如果為 True,則 tensorclass 例項中的 None 值將被寫入字典。否則,它們將被丟棄。預設值:True

  • convert_tensors (bool, "numpy") – 如果為 True,則在建立字典時,張量將被轉換為列表。如果為“numpy”,則張量將被轉換為 numpy 陣列。否則,它們將保持為張量。預設值:False

  • tolist_first (bool) – 如果為 True,則當 tensordict 具有批處理維度時,它將首先轉換為列表。預設值:False

返回:

tensordict 的字典表示。

另請參閱

tolist()

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     a=torch.arange(24).view(2, 3, 4),
...     b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)),
...     batch_size=(2, 3)
... )
>>> print(td.to_dict())
{'a': tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],
[[12, 13, 14, 15],

[16, 17, 18, 19], [20, 21, 22, 23]]]), ‘b’: {‘c’: tensor([[[ 0, 1], [ 2, 3], [ 4, 5]],

[[ 6, 7],

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

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

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

引數:
  • filename (strpath) – h5 檔案的路徑。

  • **kwargs – 傳遞給 h5py.File.create_dataset() 的其他引數。

返回:

指向新建立檔案的 PersitentTensorDict 例項。

示例

>>> import tempfile
>>> import timeit
>>>
>>> from tensordict import TensorDict, MemoryMappedTensor
>>> td = TensorDict({
...     "a": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000)),
...     "b": {"c": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000, 3))},
... }, [1_000_000])
>>>
>>> file = tempfile.NamedTemporaryFile()
>>> td_h5 = td.to_h5(file.name, compression="gzip", compression_opts=9)
>>> print(td_h5)
PersistentTensorDict(
    fields={
        a: Tensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float32, is_shared=False),
        b: PersistentTensorDict(
            fields={
                c: Tensor(shape=torch.Size([1000000, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([1000000]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([1000000]),
    device=None,
    is_shared=False)
to_lazystack(dim: int = 0)

將 TensorDict 轉換為 LazyStackedTensorDict 或等價物。

注意

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

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

dim (int, optional) – 要堆疊 tensordict 的維度。預設為 0

返回:

LazyStackedTensorDict 例項。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3))
>>> td2 = td.to_lazystack()
>>> td2.batch_size
torch.Size([2, 3])
>>> assert isinstance(td2, LazyStackedTensorDict)
to_mds(*, columns: dict[str, str] | None = None, out: str | tuple[str, str], keep_local: bool = False, compression: str | None = None, hashes: list[str] | None = None, size_limit: int | str | None = 67108864, writer: 'streaming.MDSWriter' | None = None, **kwargs) None

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

關鍵字引數:
  • out (str | Tuple[str, str]) –

    用於儲存分片檔案的輸出資料集目錄。

    1. 如果 out 是本地目錄,則分片檔案將本地儲存。

    2. 如果 out 是遠端目錄,則將建立一個本地臨時目錄來

      快取分片檔案,然後將分片檔案上傳到遠端位置。完成後,在分片上傳後刪除臨時目錄。

    3. 如果 out(local_dir, remote_dir) 的元組,則分片檔案將儲存在

      local_dir 中,並上傳到遠端位置。

  • columns (dict[str, str]) – 可選的列字典。如果未提供,將自動從張量集合中推斷。

  • keep_local (bool) – 如果資料集已上傳,是否在上傳後保留本地資料集目錄。預設為 False

  • compression (str, optional) – 可選的壓縮或壓縮級別。預設為 None

  • hashes (List[str], optional) – 可選的要應用於分片檔案的雜湊演算法列表。預設為 None

  • size_limit (Union(int, str), 可選) – 可選的分片大小限制,達到此限制後開始新分片。如果為 None,則將所有內容放入一個分片。也可以指定位元組的可讀格式,例如 "100kb" 表示 100 千位元組(100*1024)等。預設為 1 << 26。如果

  • writer – (MDSWriter, optional): 要使用的寫入器。將從 out 關鍵字引數和其他輸入關鍵字引數建立。

  • **kwargs (Any) – 寫入器的其他設定。

注意

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

警告

此方法需要安裝 mosaicml-streaming

警告

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

另請參閱

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

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

示例

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

將 TensorDictBase 例項的內容遞迴地寫入指定的 nn.Module 屬性。

to_module 也可以用作上下文管理器,用於臨時用一組引數/緩衝區填充模組(請參閱下面的示例)。

引數:

module (nn.Module) – 要將引數寫入的模組。

關鍵字引數:
  • inplace (bool, 可選) – 如果為 True,則模組中的引數或張量將原地更新。預設為 False

  • return_swap (bool, 可選) – 如果為 True,則將返回舊的引數配置。預設為 False

  • swap_dest (TensorDictBase, optional) – 如果 return_swapTrue,則為寫入交換的目標 tensordict。

  • use_state_dict (bool, 可選) – 如果為 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()
返回:

一個 tensordict,如果 return_swapTrue,則包含來自模組的值,否則為 None

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

將 tensordict 轉換為命名元組。

引數:

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

示例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> data.to_namedtuple()
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))
to_padded_tensor(padding=0.0, mask_key: Optional[NestedKey] = None) Any

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

引數:
  • padding (float) – tensordict 中張量的填充值。預設為 0.0

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

to_pytree()

將 tensordict 轉換為 PyTree。

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

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

to_struct_array() ndarray

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

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

另請參閱

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

另請參閱

呼叫 numpy() 將張量轉換為 numpy 陣列字典。

返回:

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

示例

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

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

引數:

retain_none (bool) – 如果為 True,則 tensorclass 例項中的 None 值將被寫入 tensordict。否則,它們將被丟棄。預設值:True

返回:

包含相同值的新 TensorDict 物件。

tolist(*, convert_tensors: Union[bool, Literal['numpy']] = False, tolist_first: bool = False, as_linked_list: bool = False)

tensordict.tensorclass.NonTensorStack 的內容提取到巢狀列表中。

關鍵字引數:
  • convert_tensors (bool) – 如果為 True,則張量將被轉換為列表。否則,它們將保持為張量。預設:False

  • tolist_first (bool, 可選) – 如果為 True,當 tensordict 具有批處理維度時,它將首先被轉換為列表。預設:True

  • as_linked_list (bool, 可選) – 如果為 True,列表將被轉換為 tensordict.utils.LinkedList,當列表被修改時,它將自動更新 tensordict。預設:False

示例

>>> from tensordict import NonTensorData
>>> import torch
>>> data = torch.stack([
...     torch.stack([NonTensorData(data=(i, j), batch_size=[]) for i in range(2)])
...    for j in range(3)])
>>> data.tolist()
[[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(0, 2), (1, 2)]]
transpose(dim0, dim1)

返回輸入張量的轉置張量。給定的維度 dim0dim1 將被交換。

轉置 tensordict 的原地或非原地修改也會影響原始 tensordict,因為記憶體是共享的,並且操作會映射回原始 tensordict。

示例

>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4])
>>> tensordict_transpose = tensordict.transpose(0, 1)
>>> print(tensordict_transpose.shape)
torch.Size([4, 3])
>>> tensordict_transpose.set("b",, torch.randn(4, 3))
>>> print(tensordict.get("b").shape)
torch.Size([3, 4])
trunc() Any

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

trunc_() Any

原地計算 TensorDict 每個元素的 trunc() 值。

type(dst_type: dtype) Any

將所有張量轉換為 dst_type

引數:

dst_type (type or string) – 所需型別

uint16() Any

將所有張量轉換為 torch.uint16

uint32() Any

將所有張量轉換為 torch.uint32

uint64() Any

將所有張量轉換為 torch.uint64

uint8() Any

將所有張量轉換為 torch.uint8

unbind(dim: int) tuple[T, ...]

返回一個已解綁的 tensordict 索引元組,沿指定的維度。

示例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td0, td1, td2 = td.unbind(0)
>>> td0['x']
tensor([0, 1, 2, 3])
>>> td1['x']
tensor([4, 5, 6, 7])
unflatten(dim, unflattened_size)

展平 tensordict 的一個維度,將其擴充套件為所需的形狀。

引數:
  • dim (int) – 指定要展平的輸入張量的維度。

  • unflattened_size (shape) – 是 tensordict 展平維度的新的形狀。

示例

>>> td = TensorDict({
...     "a": torch.arange(60).view(3, 4, 5),
...     "b": torch.arange(12).view(3, 4)},
...     batch_size=[3, 4])
>>> td_flat = td.flatten(0, 1)
>>> td_unflat = td_flat.unflatten(0, [3, 4])
>>> assert (td == td_unflat).all()
unflatten_keys(separator: str = '.', inplace: bool = False) Any

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

TensorDict 型別將被丟失,結果將是簡單的 TensorDict 例項。巢狀 tensordicts 的元資料將從根推斷:資料樹中的所有例項將共享相同的批次大小、維度名稱和裝置。

引數:
  • separator (str, optional) – 巢狀項之間的分隔符。

  • inplace (bool, optional) – 如果為 True,則結果 tensordict 將與呼叫它的 tensordict 具有相同的身份。預設為 False

示例

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

此方法和 unflatten_keys() 方法在處理 state-dict 時非常有用,因為它們使得能夠無縫地將扁平化字典轉換為模仿模型結構的結構。

示例

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

解鎖 tensordict 以進行非原地操作。

可用作裝飾器。

有關更多詳細資訊,請參閱 lock_()

unsqueeze(*args, **kwargs)

為介於 -td.batch_dimstd.batch_dims 之間的維度擴充套件所有張量,並將它們返回到一個新的 tensordict 中。

引數:

dim (int) – 要擴充套件的維度

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> td = td.unsqueeze(-2)
>>> td.shape
torch.Size([3, 1, 4])
>>> td.get("x").shape
torch.Size([3, 1, 4, 2])

此操作也可以用作上下文管理器。對原始 tensordict 的更改將是就地進行的,即原始張量的內容不會被修改。這也假定 tensordict 未被鎖定(否則,需要解鎖 tensordict)。

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> with td.unsqueeze(-2) as tds:
...     tds.set("y", torch.zeros(3, 1, 4))
>>> assert td.get("y").shape == [3, 4]
update(input_dict_or_td: Union[dict[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], 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, update_batch_size: bool = False, ignore_lock: bool = False) T

使用來自字典或其他 TensorDict 的值來更新 TensorDict。

警告

update 在 try/except 塊中呼叫時會損壞資料。不要在這些塊中使用此方法來嘗試捕獲和修補執行過程中發生的錯誤。

引數:
  • input_dict_or_td (TensorDictBasedict) – 要寫入 self 的輸入資料。

  • clone (bool, 可選) – 在設定輸入( tensor)字典中的張量之前是否克隆它們。預設為 False

  • inplace (bool, 可選) – 如果為 True 並且某個鍵與 tensordict 中的現有鍵匹配,則更新將對該鍵值對進行原地更新。如果找不到該條目,則會將其新增。預設為 False

關鍵字引數:
  • keys_to_update (NestedKeys 序列, 可選) – 如果提供,則僅更新 key_to_update 中的鍵列表。這旨在避免呼叫 data_dest.update(*keys_to_update))

  • non_blocking (bool, optional) – 如果為True且複製在不同裝置之間進行,則複製可能與主機非同步發生。

  • is_leaf (Callable(Type) -> bool, 可選) –

    一個指示物件型別是應被視為葉子節點並被交換,還是應被視為張量集合的可呼叫物件。

    另請參閱

    is_leaf_nontensor()default_is_leaf()

  • update_batch_size (bool, 可選) –

    如果為 True,則 update 將嘗試更新目標(self)的 batch_size,如果它與源的 batch_size 不匹配。預設為 False

    注意

    在 batch_size 不匹配的情況下,LazyStackTensorDict 例項將被清空其內容,並將使用源 tensordict 的副本重新填充容器。

    注意

    此引數假定 keys_to_update 為空,並且 inplace=False。如果目標(self)的鍵不是源鍵的子集,則會引發異常,因為 TensorDict 將無法推斷如何處理額外的目標條目。

  • ignore_lock (bool, 可選) – 如果為 True,則無論其鎖定狀態如何,都可以更新任何 tensordict。預設為 False

注意

當使用 N 個元素的 LazyStackedTensorDict 更新另一個 M 個元素的 LazyStackedTensorDict 時,沿堆疊維度,update 方法會將額外 tensordict 的副本附加到目標(self)的懶惰堆疊中。這使得使用者可以依賴 update 來逐步增加懶惰堆疊。

返回:

self

示例

>>> td = TensorDict({}, batch_size=[3])
>>> a = torch.randn(3)
>>> b = torch.randn(3, 4)
>>> other_td = TensorDict({"a": a, "b": b}, batch_size=[])
>>> td.update(other_td, inplace=True) # writes "a" and "b" even though they can't be found
>>> assert td['a'] is other_td['a']
>>> other_td = other_td.clone().zero_()
>>> td.update(other_td)
>>> assert td['a'] is not other_td['a']
update_(input_dict_or_td: Union[dict[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], T], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) T

使用來自字典或其他 TensorDict 的值原地更新 TensorDict。

update() 不同,此函式將在鍵不在 self 中時丟擲錯誤。

引數:
  • input_dict_or_td (TensorDictBasedict) – 要寫入 self 的輸入資料。

  • clone (bool, 可選) – 在設定輸入( tensor)字典中的張量之前是否克隆它們。預設為 False

關鍵字引數:
  • keys_to_update (NestedKeys 序列, 可選) – 如果提供,則僅更新 key_to_update 中的鍵列表。這旨在避免呼叫 data_dest.update_(data_src.select(*keys_to_update))

  • non_blocking (bool, optional) – 如果為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: dict[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]] | tensordict._tensorcollection.TensorCollection, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], clone: bool = False, *, non_blocking: bool = False) T

使用來自字典或其他 TensorDict 的值,在指定的索引處原地更新 TensorDict。

與 TensorDict.update 不同,此函式將在鍵未知於 TensorDict 時丟擲錯誤。

引數:
  • input_dict_or_td (TensorDictBasedict) – 要寫入 self 的輸入資料。

  • idx (int, torch.Tensor, 可迭代物件, slice) – 更新應發生的 tensordict 的索引。

  • clone (bool, 可選) – 在設定輸入( 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()
valid_keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) _LazyStackedTensorDictKeysView

返回 tensordict 鍵的生成器。

警告

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

引數:
  • include_nested (bool, 可選) – 如果為 `True`,則返回巢狀值。預設為 `False`。

  • leaves_only (bool, 可選) – 如果為 `False`,則僅返回葉子節點。預設為 `False`。

  • is_leaf (callable, optional) –

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

    注意

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

    另請參閱

    is_leaf_nontensor()default_is_leaf()

關鍵字引數:

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

示例

>>> from tensordict import TensorDict
>>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[])
>>> data.keys()
['0', '1']
>>> list(data.keys(leaves_only=True))
['0']
>>> list(data.keys(include_nested=True, leaves_only=True))
['0', '1', ('1', '2')]
values(include_nested=False, leaves_only=False, is_leaf=None, *, sort: bool = False)

返回一個表示 tensordict 值的生成器。

引數:
  • include_nested (bool, 可選) – 如果為 `True`,則返回巢狀值。預設為 `False`。

  • leaves_only (bool, 可選) – 如果為 `False`,則僅返回葉子節點。預設為 `False`。

  • is_leaf (callable, optional) –

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

    注意

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

    另請參閱

    is_leaf_nontensor()default_is_leaf()

關鍵字引數:

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

var(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) Union[Any, Tensor]

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

引數:
  • dim (int, int 元組, 可選) – 如果為 None,則返回一個無量綱的 tensordict,其中包含所有葉子的總和值(如果可以計算)。如果為整數或整數元組,則當此維度與 tensordict 形狀相容時,將在指定的維度上呼叫 var。當前只允許 “feature” 字串。使用 dim=”feature” 將在所有特徵維度上執行歸約。如果 reduce=True,則將返回形狀與 TensorDict 的 batch_size 相同的張量。否則,將返回一個結構與 self 相同的、特徵維度已歸約的新 tensordict。

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

關鍵字引數:
  • correction (int) – 樣本大小與樣本自由度之差。預設為貝塞爾校正,correction=1。

  • reduce (bool, optional) – 如果 True,則會跨所有 TensorDict 值進行歸約,並返回單個歸約後的張量。預設為 False

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.var(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.var()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.var(reduce=True)
tensor(1.0006)
>>> td.var(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.var(reduce=True, dim="feature")
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
>>> td.var(reduce=True, dim=0)
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])
view(*shape: int, size: list | tuple | torch.Size | None = None, batch_size: torch.Size | None = None)

返回一個 tensordict,其中包含與 tensordict batch_size 相容的新形狀的張量檢視。

或者,可以將 dtype 作為第一個未命名引數提供。在這種情況下,所有張量都將以相應的 dtype 進行檢視處理。請注意,這假定新的形狀與提供的 dtype 相容。有關 dtype 檢視的更多資訊,請參閱 view()

引數:
  • *shape (int) – 結果 TensorDict 的新形狀。

  • dtype (torch.dtype) – 或者,用於表示張量內容的 dtype。

  • size – 可迭代物件

關鍵字引數:

batch_size (torch.Size, 可選) – 如果提供了 dtype,則可以使用此關鍵字引數重置 batch_size。如果 view 是使用形狀呼叫的,則此引數無效。

返回:

具有所需 batch_size 的新 tensordict。

示例

>>> td = TensorDict(source={'a': torch.zeros(3,4,5),
...    'b': torch.zeros(3,4,10,1)}, batch_size=torch.Size([3, 4]))
>>> td_view = td.view(12)
>>> print(td_view.get("a").shape)  # torch.Size([12, 5])
>>> print(td_view.get("b").shape)  # torch.Size([12, 10, 1])
>>> td_view = td.view(-1, 4, 3)
>>> print(td_view.get("a").shape)  # torch.Size([1, 4, 3, 5])
>>> print(td_view.get("b").shape)  # torch.Size([1, 4, 3, 10, 1])
where(condition: Tensor, other: torch.Tensor | tensordict.base.TensorDictBase, *, out: Optional[TensorDictBase] = None, pad: Optional[Union[int, bool]] = None, update_batch_size: bool = False)

根據 condition 從 self 或 other 中選擇元素,並返回一個 TensorDict

引數:
  • condition (BoolTensor) – 當 True (非零) 時,返回 self,否則返回 other

  • other (TensorDictBase標量) – 當 other 是標量時,值為標量;或者,當 conditionFalse 時的值為選定的值。

關鍵字引數:
  • out (TensorDictBase, 可選) – 輸出 TensorDictBase 例項。

  • pad (標量, 可選) – 如果提供,源或目標 tensordict 中缺失的鍵將寫入為 torch.where(mask, self, pad)torch.where(mask, pad, other)。預設為 None,即不允許缺失的鍵。

  • update_batch_size (bool, optional) – 如果 Trueout 已提供,則輸出的批次大小將更新以匹配 condition 的批次大小。預設為 False

zero_() Any

原地將 tensordict 中的所有張量清零。

zero_grad(set_to_none: bool = True) Any

遞迴地將 TensorDict 的所有梯度清零。

引數:

set_to_none (bool, 可選) – 如果為 True,則 tensor.grad 將為 None,否則為 0。預設為 True

文件

訪問全面的 PyTorch 開發者文件

檢視文件

教程

為初學者和高階開發者提供深入的教程

檢視教程

資源

查詢開發資源並讓您的問題得到解答

檢視資源