TensorDictBase¶
- class tensordict.TensorDictBase¶
TensorDictBase 是 TensorDict 的抽象父類,TensorDict 是 torch.Tensor 資料容器。
- abs() T¶
計算 TensorDict 中每個元素的絕對值。
- abs_() T¶
就地計算 TensorDict 中每個元素的絕對值。
- acos() T¶
計算 TensorDict 中每個元素的
acos()值。
- acos_() T¶
就地計算 TensorDict 中每個元素的
acos()值。
- add(other: tensordict.base.TensorDictBase | torch.Tensor, *, alpha: float | None = None, default: str | torch.Tensor | None = None) TensorDictBase¶
將
other乘以alpha後加到self。\[\text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i\]- 引數:
other (TensorDictBase 或 torch.Tensor) – 要新增到
self的張量或 TensorDict。- 關鍵字引數:
alpha (Number, optional) –
other的乘數。default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- add_(other: tensordict.base.TensorDictBase | torch.Tensor | float, *, alpha: Optional[float] = None)¶
add()的原地版本。注意
原地
add不支援default關鍵字引數。
- addcdiv(other1: tensordict.base.TensorDictBase | torch.Tensor, other2: tensordict.base.TensorDictBase | torch.Tensor, value: float | None = 1)¶
執行
other1除以other2的逐元素除法,將結果乘以標量value,然後將其加到self中。\[\text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i}\]self、other1和other2的元素的形狀必須是可廣播的。對於 FloatTensor 或 DoubleTensor 型別的輸入,
value必須是實數,否則必須是整數。- 引數:
other1 (TensorDict 或 Tensor) – 分子 tensordict(或 tensor)
tensor2 (TensorDict 或 Tensor) – 分母 tensordict(或 tensor)
- 關鍵字引數:
value (Number, optional) – \(\text{tensor1} / \text{tensor2}\) 的乘數
- addcmul(other1: tensordict.base.TensorDictBase | torch.Tensor, other2: tensordict.base.TensorDictBase | torch.Tensor, *, value: float | None = 1)¶
執行
other1乘以other2的逐元素乘法,將結果乘以標量value,然後將其加到self中。\[\text{out}_i = \text{input}_i + \text{value} \times \text{other1}_i \times \text{other2}_i\]self、other1和other2的形狀必須是可廣播的。對於 FloatTensor 或 DoubleTensor 型別的輸入,
value必須是實數,否則必須是整數。- 引數:
other1 (TensorDict 或 Tensor) – 要相乘的 tensordict 或 tensor
other2 (TensorDict 或 Tensor) – 要相乘的 tensordict 或 tensor
- 關鍵字引數:
value (Number, optional) – \(other1 .* other2\) 的乘數
- abstract all(dim: Optional[int] = None) bool | tensordict.base.TensorDictBase¶
檢查 tensordict 中所有值是否為 True/非空。
- 引數:
dim (int, optional) – 如果
None,則返回一個布林值,指示是否所有 tensor 都返回 tensor.all() == True。如果為整數,則僅當此維度與 tensordict 形狀相容時,才對指定的維度呼叫 all。
- amax(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None) TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的最大值。
與
max()相同,但return_indices=False。
- amin(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None) TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的最小值。
與
min()相同,但return_indices=False。
- abstract any(dim: Optional[int] = None) bool | tensordict.base.TensorDictBase¶
檢查 tensordict 中是否有任何值為 True/非空。
- 引數:
dim (int, optional) – 如果
None,則返回一個布林值,指示是否所有張量都返回 tensor.any() == True。 如果是整數,則僅當此維度與 tensordict 形狀相容時,才對指定的維度呼叫 all。
- apply(fn: Callable, *others: T, batch_size: Optional[Sequence[int]] = None, device: torch.device | None = _NoDefault.ZERO, names: Optional[Sequence[str]] = _NoDefault.ZERO, inplace: bool = False, default: Any = _NoDefault.ZERO, filter_empty: Optional[bool] = None, propagate_lock: bool = False, call_on_nested: bool = False, out: Optional[TensorDictBase] = None, **constructor_kwargs) Optional[T]¶
將可呼叫物件應用於儲存在 tensordict 中的所有值,並將它們設定在一個新的 tensordict 中。
可呼叫物件的簽名必須是
Callable[Tuple[Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。- 引數:
fn (Callable) – 要應用於 tensordict 中張量的函式。
*others (TensorDictBase 例項, 可選) – 如果提供,這些 tensordict 例項應具有與 self 相同的結構。
fn引數應接收與 tensordict 數量一樣多的未命名輸入,包括 self。 如果其他 tensordict 缺少條目,可以透過default關鍵字引數傳遞預設值。
- 關鍵字引數:
batch_size (int 序列, 可選) – 如果提供,則生成的 TensorDict 將具有所需的 batch_size。
batch_size引數應與轉換後的 batch_size 相匹配。 這是一個僅限關鍵字的引數。device (torch.device, 可選) – 結果裝置,如果有。
names (str 列表, 可選) – 新的維度名稱,以防 batch_size 被修改。
inplace (bool, 可選) – 如果為 True,則進行就地更改。 預設為 False。 這是一個僅限關鍵字的引數。
default (Any, 可選) – 其他 tensordict 中缺失條目的預設值。 如果未提供,則缺失條目將引發 KeyError。
filter_empty (bool, 可選) – 如果為
True,則將過濾掉空的 tensordict。 這也帶來了更低的計算成本,因為不會建立和銷燬空資料結構。 非張量資料被視為葉節點,因此即使該函式未觸及,也會保留在 tensordict 中。 為了向後相容,預設為False。propagate_lock (bool, 可選) – 如果為
True,則鎖定的 tensordict 將生成另一個鎖定的 tensordict。 預設為False。call_on_nested (bool, 可選) –
如果
True,該函式將被呼叫在第一層級的張量和容器(TensorDict 或 tensorclass)上。在這種情況下,func負責將其呼叫傳播到巢狀的層級。這允許在將呼叫傳播到巢狀的 tensordict 時進行細粒度的控制。如果False,該函式將只被呼叫在葉節點上,並且apply將負責將該函式分派到所有葉節點。>>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]}) >>> def mean_tensor_only(val): ... if is_tensor_collection(val): ... raise RuntimeError("Unexpected!") ... return val.mean() >>> td_mean = td.apply(mean_tensor_only) >>> def mean_any(val): ... if is_tensor_collection(val): ... # Recurse ... return val.apply(mean_any, call_on_nested=True) ... return val.mean() >>> td_mean = td.apply(mean_any, call_on_nested=True)
out (TensorDictBase, 可選) –
一個用於寫入結果的 tensordict。這可以用於避免建立新的 tensordict
>>> td = TensorDict({"a": 0}) >>> td.apply(lambda x: x+1, out=td) >>> assert (td==1).all()
警告
如果在 tensordict 上執行的操作需要訪問多個鍵才能進行單個計算,則提供一個等於
self的out引數可能會導致該操作提供靜默的錯誤結果。例如:>>> td = TensorDict({"a": 1, "b": 1}) >>> td.apply(lambda x: x+td["a"])["b"] # Right! tensor(2) >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong! tensor(3)
**constructor_kwargs – 要傳遞給 TensorDict 建構函式的其他關鍵字引數。
- 返回:
一個新的 tensordict,包含 transformed_in 張量。
示例
>>> td = TensorDict({ ... "a": -torch.ones(3), ... "b": {"c": torch.ones(3)}}, ... batch_size=[3]) >>> td_1 = td.apply(lambda x: x+1) >>> assert (td_1["a"] == 0).all() >>> assert (td_1["b", "c"] == 2).all() >>> td_2 = td.apply(lambda x, y: x+y, td) >>> assert (td_2["a"] == -2).all() >>> assert (td_2["b", "c"] == 2).all()
注意
如果函式返回
None,則該條目將被忽略。這可用於過濾 tensordict 中的資料>>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, []) >>> def filter(tensor): ... if tensor == 1: ... return tensor >>> td.apply(filter) TensorDict( fields={ 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
注意
apply 方法將返回一個
TensorDict例項,而與輸入型別無關。要保持相同的型別,可以執行>>> out = td.clone(False).update(td.apply(...))
- apply_(fn: Callable, *others, **kwargs) T¶
將一個可呼叫物件應用於 tensordict 中儲存的所有值,並將它們重新寫入原地。
- 引數:
fn (Callable) – 要應用於 tensordict 中張量的函式。
*others (TensorDictBase 序列, 可選) – 要使用的其他 tensordict。
關鍵字引數: 請參閱
apply()。- 返回:
self 或應用了函式的 self 的副本
- asin() T¶
計算 TensorDict 中每個元素的
asin()值。
- asin_() T¶
計算 TensorDict 中每個元素的
asin()值,並原地修改。
- atan() T¶
計算 TensorDict 中每個元素的
atan()值。
- atan_() T¶
計算 TensorDict 中每個元素的
atan()值,並原地修改。
- auto_batch_size_(batch_dims: Optional[int] = None) T¶
設定 tensordict 的最大批次大小,最多到可選的 batch_dims。
- 引數:
batch_dims (int, 可選) – 如果提供,則批次大小最多為
batch_dims長。- 返回:
self
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict({"a": torch.randn(3, 4, 5), "b": {"c": torch.randn(3, 4, 6)}}, batch_size=[]) >>> td.auto_batch_size_() >>> print(td.batch_size) torch.Size([3, 4]) >>> td.auto_batch_size_(batch_dims=1) >>> print(td.batch_size) torch.Size([3])
- auto_device_() T¶
自動設定裝置,如果它是唯一的。
返回: 帶有已編輯的
device屬性的 self。
- abstract property batch_size: Size¶
TensorDict 的形狀(或批次大小)。
tensordict 的形狀對應於它包含的張量的共同的第一個
N維度,其中N是一個任意數字。批次大小與 “特徵大小” 形成對比,“特徵大小” 表示張量的語義相關的形狀。例如,一批影片的形狀可能是[B, T, C, W, H],其中[B, T]是批次大小(批次和時間維度),[C, W, H]是特徵維度(通道和空間維度)。TensorDict的形狀由使用者在初始化時控制(即,它不是從張量形狀推斷出來的)。如果新大小與 TensorDict 內容相容,則可以動態編輯
batch_size。例如,始終允許將批次大小設定為空值。- 返回:
描述 TensorDict 批次大小的
Size物件。
示例
>>> data = TensorDict({ ... "key 0": torch.randn(3, 4), ... "key 1": torch.randn(3, 5), ... "nested": TensorDict({"key 0": torch.randn(3, 4)}, batch_size=[3, 4])}, ... batch_size=[3]) >>> data.batch_size = () # resets the batch-size to an empty value
- bfloat16()¶
將所有張量轉換為
torch.bfloat16。
- bitwise_and(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) TensorDictBase¶
執行
self和other之間的按位與操作。\[\text{{out}}_i = \text{{input}}_i \land \text{{other}}_i\]- 引數:
other (TensorDictBase 或 torch.Tensor) – 用於執行按位與操作的張量或 TensorDict。
- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- bool()¶
將所有張量轉換為
torch.bool。
- bytes(*, count_duplicates: bool = True) int¶
計算包含的張量的位元組數。
- 關鍵字引數:
count_duplicates (bool) – 是否將重複的張量計數為獨立的張量。 如果
False,則僅會丟棄嚴格相同的張量(來自公共基礎張量的相同檢視但不同的 id 將被計數兩次)。 預設為 True (假設每個張量都是一個單獨的副本)。
- classmethod cat(input, dim=0, *, out=None)¶
沿給定維度將多個 tensordict 串聯成一個 tensordict。
此呼叫等效於呼叫
torch.cat(),但與 torch.compile 相容。
- cat_from_tensordict(dim: int = 0, *, sorted: Optional[Union[bool, List[NestedKey]]] = None, out: Optional[Tensor] = None) Tensor¶
將 tensordict 的所有條目串聯成一個張量。
- 引數:
dim (int, optional) – 條目應沿其串聯的維度。
- 關鍵字引數:
sorted (bool 或 NestedKey 列表) – 如果
True,則條目將按字母順序串聯。 如果為False(預設),將使用 dict 順序。 或者,可以提供一個鍵名列表,張量將相應地進行串聯。 這會產生一些開銷,因為將根據 tensordict 中的葉名稱列表檢查鍵列表。out (torch.Tensor, optional) – cat 操作的可選目標張量。
- cat_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) T¶
將條目串聯到一個新條目中,並可能刪除原始值。
- 引數:
keys (NestedKey 序列) – 要串聯的條目。
- 關鍵字引數
out_key (NestedKey): 串聯輸入的新鍵名。 keep_entries (bool, optional): 如果
False,keys中的條目將被刪除。預設為
False。- dim (int, optional): 必須發生串聯的維度。
預設為
0。
返回值: self
示例
>>> td = TensorDict(a=torch.zeros(1), b=torch.ones(1)) >>> td.cat_tensors("a", "b", out_key="c") >>> assert "a" not in td >>> assert (td["c"] == torch.tensor([0, 1])).all()
- ceil() T¶
計算 TensorDict 中每個元素的
ceil()值。
- ceil_() T¶
就地計算 TensorDict 中每個元素的
ceil()值。
- chunk(chunks: int, dim: int = 0) tuple[tensordict.base.TensorDictBase, ...]¶
如果可能,將 tensordict 拆分為指定數量的塊。
每個塊都是輸入 tensordict 的一個檢視。
示例
>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 4, 2), ... }, batch_size=[3, 4]) >>> td0, td1 = td.chunk(dim=-1, chunks=2) >>> td0['x'] tensor([[[ 0, 1], [ 2, 3]], [[ 8, 9], [10, 11]], [[16, 17], [18, 19]]])
- clamp(min: tensordict.base.TensorDictBase | torch.Tensor = None, max: tensordict.base.TensorDictBase | torch.Tensor = None, *, out=None)¶
將
self中的所有元素鉗制到範圍 [min,max] 中。令 min_value 和 max_value 分別為
min和max,這將返回注意
如果
min大於maxtorch.clamp(..., min, max)將input中的所有元素設定為max的值。
- clamp_max(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) T¶
將
self的元素限制到other,如果它們大於該值。- 引數:
other (TensorDict 或者 Tensor) – 另一個輸入的 tensordict 或 tensor。
- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- clamp_max_(other: tensordict.base.TensorDictBase | torch.Tensor) T¶
clamp_max()的原地版本。注意
原地
clamp_max不支援default關鍵字引數。
- clamp_min(other: tensordict.base.TensorDictBase | torch.Tensor, default: str | torch.Tensor | None = None) T¶
將
self的元素限制到other,如果它們小於該值。- 引數:
other (TensorDict 或者 Tensor) – 另一個輸入的 tensordict 或 tensor。
- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- clamp_min_(other: tensordict.base.TensorDictBase | torch.Tensor) T¶
clamp_min()的原地版本。注意
原地
clamp_min不支援default關鍵字引數。
- clear() T¶
清除 tensordict 的內容。
- clear_device_() T¶
清除 tensordict 的裝置。
返回值: self
- clear_refs_for_compile_() T¶
清除 weakrefs,以便 tensordict 安全地離開編譯區域。
在返回 TensorDict 之前,如果您遇到 torch._dynamo.exc.Unsupported: reconstruct: WeakRefVariable(),請使用此方法。
返回值: self
- clone(recurse: bool = True, **kwargs) T¶
將 TensorDictBase 子類的例項克隆到相同型別的新 TensorDictBase 子類上。
要從任何其他 TensorDictBase 子型別建立 TensorDict 例項,請呼叫
to_tensordict()方法。- 引數:
recurse (bool, optional) – 如果
True,TensorDict 中包含的每個 tensor 也將被複制。否則,只會複製 TensorDict 樹結構。預設為True。
注意
與許多其他操作(逐點算術、形狀操作等)不同,
clone不繼承原始鎖屬性。 這種設計選擇的目的是為了建立一個可以修改的克隆,這是最常見的用法。
- complex128()¶
將所有 tensor 轉換為
torch.complex128。
- complex32()¶
將所有 tensor 轉換為
torch.complex32。
- complex64()¶
將所有 tensor 轉換為
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, 可選) – 用於記憶體對映張量的可選檔案路徑,用作 tensordict 的儲存。
- 關鍵字引數:
num_threads (integer, 可選) – 用於填充儲存的執行緒數。
device (torch.device, 可選) – 儲存必須例項化的可選裝置。
inplace (bool, 可選) – 如果
True,則生成的 tensordict 與self相同,但值已更新。預設為False。return_early (bool, 可選) – 如果
True且num_threads>0,則該方法將返回 tensordict 的 future。可以使用 future.result() 查詢生成的 tensordict。use_buffer (bool, 可選) – 如果
True且傳遞了檔名,則將在共享記憶體中建立一箇中間本地緩衝區,並且資料將作為最後一步複製到儲存位置。這可能比直接寫入遠端物理記憶體(例如,NFS)更快。預設為False。share_memory (bool, 可選) – 如果
True,儲存將放置在共享記憶體中。預設為False。pin_memory (bool, 可選) – 整合的資料是否應放置在固定記憶體中。預設為
False。metadata (bool, 可選) – 如果
True,元資料將與公共儲存一起儲存。如果提供了檔名,則此操作無效。當人們想要控制如何實現序列化時,儲存元資料可能很有用,因為如果元資料可用或不可用,TensorDict 會以不同的方式處理整合的 TD 的 pickle/unpickle。
注意
如果 tensordict 已經整合,則所有引數都將被忽略,並返回
self。呼叫contiguous()重新整合。示例
>>> import pickle >>> import tempfile >>> import torch >>> import tqdm >>> from torch.utils.benchmark import Timer >>> from tensordict import TensorDict >>> data = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}}) >>> data_consolidated = data.consolidate() >>> # check that the data has a single data_ptr() >>> assert torch.tensor([ ... v.untyped_storage().data_ptr() for v in data_c.values(True, True) ... ]).unique().numel() == 1 >>> # Serializing the tensordict will be faster with data_consolidated >>> with open("data.pickle", "wb") as f: ... print("regular", Timer("pickle.dump(data, f)", globals=globals()).adaptive_autorange()) >>> with open("data_c.pickle", "wb") as f: ... print("consolidated", Timer("pickle.dump(data_consolidated, f)", globals=globals()).adaptive_autorange())
- abstract contiguous() T¶
返回一個相同型別的新 tensordict,其值為連續的(如果值已經是連續的,則返回 self)。
- copy()¶
返回 tensordict 的淺複製(即,複製結構但不復制資料)。
等效於 TensorDictBase.clone(recurse=False)
- copy_(tensordict: T, non_blocking: bool = False) T¶
-
non-blocking 引數將被忽略,僅用於與
torch.Tensor.copy_()相容。
- copy_at_(tensordict: T, idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], non_blocking: bool = False) T¶
- cos() T¶
計算 TensorDict 中每個元素的
cos()值。
- cos_() T¶
原地計算 TensorDict 中每個元素的
cos()值。
- cosh() T¶
計算 TensorDict 中每個元素的
cosh()值。
- cosh_() T¶
原地計算 TensorDict 中每個元素的
cosh()值。
- create_nested(key)¶
建立一個與當前 tensordict 具有相同形狀、裝置和維度名稱的巢狀 tensordict。
如果該值已經存在,此操作將會覆蓋它。 此操作在鎖定的 tensordict 中會被阻塞。
示例
>>> data = TensorDict({}, [3, 4, 5]) >>> data.create_nested("root") >>> data.create_nested(("some", "nested", "value")) >>> print(data) TensorDict( fields={ root: TensorDict( fields={ }, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False), some: TensorDict( fields={ nested: TensorDict( fields={ value: TensorDict( fields={ }, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)
- cuda(device: Optional[int] = None, **kwargs) T¶
將 tensordict 轉換為 cuda 裝置(如果尚未在其上)。
- 引數:
device (int, 可選) – 如果提供,則指定 tensor 應該轉換到的 cuda 裝置。
該函式還支援
to()的所有關鍵字引數。
- cummax(dim: int, *, reduce: Optional[bool] = None, return_indices: bool = True) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的累積最大值。
- 引數:
dim (int) – 整數,表示執行 cummax 操作的維度。
- 關鍵字引數:
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.cummax(dim=0) cummax( indices=TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False), vals=TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False), batch_size=torch.Size([4]), device=None, is_shared=False) >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.cummax(reduce=True, dim=0) torch.return_types.cummax(...)
- cummin(dim: int, *, reduce: Optional[bool] = None, return_indices: bool = True) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的累積最小值。
- 引數:
dim (int) – 表示執行 cummin 操作的維度的整數。
- 關鍵字引數:
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.cummin(dim=0) cummin( indices=TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False), vals=TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False), batch_size=torch.Size([4]), device=None, is_shared=False) >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.cummin(reduce=True, dim=0) torch.return_types.cummin(...)
- property data¶
返回一個 tensordict,其中包含葉張量的 .data 屬性。
- data_ptr(*, storage: bool = False)¶
返回 tensordict 葉子的 data_ptr。
這對於檢查兩個 tensordict 是否共享相同的
data_ptr()很有用。- 關鍵字引數:
storage (bool, optional) – 如果為
True,則會呼叫 tensor.untyped_storage().data_ptr()。 預設為False。
示例
>>> from tensordict import TensorDict >>> td = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2]) >>> assert (td0.data_ptr() == td.data_ptr()).all()
注意
LazyStackedTensorDict例項將顯示為巢狀的 tensordict,以反映其葉子的真實data_ptr()>>> td0 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2]) >>> td1 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2]) >>> td = TensorDict.lazy_stack([td0, td1]) >>> td.data_ptr() TensorDict( fields={ 0: TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False), 1: TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False)
- abstract del_(key: NestedKey) T¶
刪除 tensordict 的一個鍵。
- 引數:
key (NestedKey) – 要刪除的鍵
- 返回:
self
- densify(layout: layout = torch.strided)¶
嘗試用連續張量(普通張量或巢狀張量)表示延遲堆疊。
- 關鍵字引數:
layout (torch.layout) – 巢狀張量的佈局(如果有)。 預設為
strided。
- detach() T¶
分離 tensordict 中的張量。
- 返回:
一個新的 tensordict,其中沒有張量需要梯度。
- abstract detach_() T¶
就地分離 tensordict 中的張量。
- 返回:
self。
- abstract 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: str | torch.Tensor | None = None) T¶
將輸入
self的每個元素除以other的相應元素。\[\text{out}_i = \frac{\text{input}_i}{\text{other}_i}\]支援廣播、型別提升以及整數、浮點數、tensordict 或張量輸入。 始終將整數型別提升為預設標量型別。
- 引數:
other (TensorDict, Tensor or Number) – 除數。
- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- div_(other: tensordict.base.TensorDictBase | torch.Tensor) T¶
div()的原地 (in-place) 版本。注意
原地
div不支援default關鍵字引數。
- double()¶
將所有張量轉換為
torch.bool。
- property dtype¶
返回 tensordict 中值的 dtype,如果它是唯一的。
- dumps(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T¶
將 tensordict 儲存到磁碟。
此函式是
memmap()的代理。
- empty(recurse=False, *, batch_size=None, device=_NoDefault.ZERO, names=None) T¶
返回一個新的、空的 tensordict,具有相同的裝置和批次大小。
- 引數:
recurse (bool, 可選) – 如果
True,則會複製TensorDict的整個結構,但不包含內容。 否則,只會複製根。 預設為False。- 關鍵字引數:
batch_size (torch.Size, 可選) – tensordict 的新批次大小。
device (torch.device, 可選) – 新裝置。
names (字串列表, 可選) – 維度名稱。
- abstract entry_class(key: NestedKey) type¶
返回條目的類,可能避免呼叫 isinstance(td.get(key), type)。
只要
get()的執行成本很高,就應該優先使用此方法而不是tensordict.get(key).shape。
- erf() T¶
計算 TensorDict 中每個元素的
erf()值。
- erf_() T¶
原地計算 TensorDict 中每個元素的
erf()值。
- erfc() T¶
計算 TensorDict 中每個元素的
erfc()值。
- erfc_() T¶
原地計算 TensorDict 中每個元素的
erfc()值。
- exclude(*keys: NestedKey, inplace: bool = False) T¶
排除 tensordict 的鍵,並返回一個不包含這些條目的新 tensordict。
這些值未被複制:對原始或新 tensordict 中的張量進行原地修改會導致兩個 tensordict 都發生更改。
- 引數:
- 返回:
沒有排除條目的新 tensordict(如果
inplace=True,則相同)。
示例
>>> from tensordict import TensorDict >>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, []) >>> td.exclude("a", ("b", "c")) TensorDict( fields={ b: TensorDict( fields={ d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.exclude("a", "b") TensorDict( fields={ }, batch_size=torch.Size([]), device=None, is_shared=False)
- exp() T¶
計算 TensorDict 中每個元素的
exp()值。
- exp_() T¶
原地計算 TensorDict 中每個元素的
exp()值。
- abstract expand(*shape: int) T¶
- abstract expand(shape: Size) T
根據
expand()函式擴充套件 tensordict 中的每個張量,忽略特徵維度。支援使用可迭代物件來指定形狀。
示例
>>> td = TensorDict({ ... 'a': torch.zeros(3, 4, 5), ... 'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4]) >>> td_expand = td.expand(10, 3, 4) >>> assert td_expand.shape == torch.Size([10, 3, 4]) >>> assert td_expand.get("a").shape == torch.Size([10, 3, 4, 5])
- expand_as(other: tensordict.base.TensorDictBase | torch.Tensor) TensorDictBase¶
將 tensordict 的形狀廣播到 other 的形狀,並據此進行擴充套件。
如果輸入是張量集合(tensordict 或 tensorclass),則葉子節點將以一對一的方式進行擴充套件。
示例
>>> from tensordict import TensorDict >>> import torch >>> td0 = TensorDict({ ... "a": torch.ones(3, 1, 4), ... "b": {"c": torch.ones(3, 2, 1, 4)}}, ... batch_size=[3], ... ) >>> td1 = TensorDict({ ... "a": torch.zeros(2, 3, 5, 4), ... "b": {"c": torch.zeros(2, 3, 2, 6, 4)}}, ... batch_size=[2, 3], ... ) >>> expanded = td0.expand_as(td1) >>> assert (expanded==1).all() >>> print(expanded) TensorDict( fields={ a: Tensor(shape=torch.Size([2, 3, 5, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([2, 3, 2, 6, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([2, 3]), device=None, is_shared=False)}, batch_size=torch.Size([2, 3]), device=None, is_shared=False)
- expm1() T¶
計算 TensorDict 中每個元素的
expm1()值。
- expm1_() T¶
就地計算 TensorDict 中每個元素的
expm1()值。
- filter_empty_()¶
就地過濾掉所有空的 tensordict。
- filter_non_tensor_data() T¶
過濾掉所有非張量資料。
- flatten(start_dim=0, end_dim=- 1)¶
展平 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_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) T¶
將巢狀的 tensordict 遞迴地轉換為扁平的 tensordict。
TensorDict 型別將丟失,結果將是一個簡單的 TensorDict 例項。
- 引數:
separator (str, 可選) – 巢狀項之間的分隔符。
inplace (bool, 可選) – 如果為
True,則生成的 tensordict 將與呼叫它的 tensordict 具有相同的標識。預設為False。is_leaf (callable, 可選) –
一個對類型別進行呼叫的可呼叫物件,返回一個布林值,指示該類是否必須被視為葉子節點。
注意
is_leaf 的目的不是阻止對巢狀 tensordict 的遞迴呼叫,而是將某些型別標記為 “葉子”,以便在 leaves_only=True 時進行過濾。即使 is_leaf(cls) 返回 True,如果 include_nested=True,tensordict 的巢狀結構仍將被遍歷。換句話說,is_leaf 不控制遞迴深度,而是提供一種方法,當 leaves_only=True 時,從結果中過濾掉某些型別。這意味著樹中的節點既可以是葉子節點,也可以是具有子節點的節點。實際上,
is_leaf的預設值確實將 tensordict 和 tensorclass 例項從葉子集合中排除。
示例
>>> data = TensorDict({"a": 1, ("b", "c"): 2, ("e", "f", "g"): 3}, batch_size=[]) >>> data.flatten_keys(separator=" - ") TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b - c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), e - f - g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
此方法和
unflatten_keys()在處理 state-dicts 時特別有用,因為它們可以無縫地將扁平字典轉換為模仿模型結構的資料結構。示例
>>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4)) >>> ddp_model = torch.ao.quantization.QuantWrapper(model) >>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".") >>> print(state_dict) TensorDict( fields={ module: TensorDict( fields={ 0: TensorDict( fields={ bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> model_state_dict = state_dict.get("module") >>> print(model_state_dict) TensorDict( fields={ 0: TensorDict( fields={ bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> model.load_state_dict(dict(model_state_dict.flatten_keys(".")))
- float()¶
將所有張量轉換為
torch.float。
- float16()¶
將所有張量轉換為
torch.float16。
- float32()¶
將所有張量轉換為
torch.float32。
- float64()¶
將所有張量轉換為
torch.float64。
- floor() T¶
計算 TensorDict 中每個元素的
floor()值。
- floor_() T¶
就地計算 TensorDict 中每個元素的
floor()值。
- frac() T¶
計算 TensorDict 中每個元素的
frac()值。
- frac_() T¶
原地計算 TensorDict 中每個元素的
frac()值。
- classmethod from_any(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶
遞迴地將任何物件轉換為 TensorDict。
注意
from_any比常規 TensorDict 建構函式限制更少。它可以使用自定義啟發式方法將 dataclass 或 tuple 等資料結構轉換為 tensordict。 這種方法可能會產生一些額外的開銷,並且涉及在對映策略方面更加主觀的選擇。注意
此方法遞迴地將輸入物件轉換為 TensorDict。如果該物件已經是 TensorDict(或任何類似的 tensor 集合物件),它將按原樣返回。
- 引數:
obj – 要轉換的物件。
- 關鍵字引數:
auto_batch_size (bool, 可選) – 如果
True,將自動計算批次大小。預設為False。batch_dims (int, 可選) – 如果 auto_batch_size 為
True,則定義輸出 tensordict 應具有的維度數。 預設為None(每個級別的完整批次大小)。device (torch.device, 可選) – 將在其上建立 TensorDict 的裝置。
batch_size (torch.Size, 可選) – TensorDict 的批次大小。 與
auto_batch_size互斥。
- 返回:
輸入物件的 TensorDict 表示。
支援的物件
透過
from_dataclass()的 Dataclass(dataclass 將被轉換為 TensorDict 例項,而不是 tensorclass)。透過
from_namedtuple()的 Namedtuple。透過
from_dict()的字典。透過
from_tuple()的元組。透過
from_struct_array()的 NumPy 的結構化陣列。透過
from_h5()的 HDF5 物件。
- classmethod from_dataclass(dataclass, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, as_tensorclass: bool = False, device: Optional[device] = None, batch_size: Optional[Size] = None)¶
將 dataclass 轉換為 TensorDict 例項。
- 引數:
dataclass – 要轉換的 dataclass 例項。
- 關鍵字引數:
auto_batch_size (bool, 可選) – 如果
True,自動確定並將批次大小應用於生成的 TensorDict。預設為False。batch_dims (int, 可選) – 如果
auto_batch_size為True,則定義輸出 tensordict 應具有的維度數。 預設為None(每個級別的完整批次大小)。as_tensorclass (bool, 可選) – 如果
True,將轉換委託給自由函式from_dataclass()並返回一個 tensor 相容的類 (tensorclass()) 或例項,而不是 TensorDict。 預設為False。device (torch.device, 可選) – 將在其上建立 TensorDict 的裝置。 預設為
None。batch_size (torch.Size, 可選) – TensorDict 的批次大小。 預設為
None。
- 返回:
從提供的 dataclass 派生的 TensorDict 例項,除非 as_tensorclass 為 True,在這種情況下,將返回一個 tensor 相容的類或例項。
- 引發:
TypeError – 如果提供的輸入不是 dataclass 例項。
警告
此方法不同於自由函式 from_dataclass 並且用途不同。 雖然自由函式返回 tensor 相容的類或例項,但此方法返回 TensorDict 例項。
- abstract classmethod from_dict(input_dict, *, auto_batch_size: Optional[bool] = None, batch_size: Optional[Size] = None, device: Optional[device] = None, batch_dims: Optional[int] = None, names: Optional[List[str]] = None)¶
從字典或另一個
TensorDict建立一個 TensorDict 並返回。如果未指定
batch_size,則返回可能的最大批次大小。此函式也適用於巢狀字典,或者可用於確定巢狀 tensordict 的批次大小。
- 引數:
input_dict (字典, 可選) – 用作資料來源的字典(相容巢狀鍵)。
- 關鍵字引數:
auto_batch_size (bool, 可選) – 如果
True,將自動計算批次大小。預設為False。batch_size (int 的可迭代物件, 可選) – tensordict 的批次大小。
device (torch.device 或 相容型別, 可選) – TensorDict 的裝置。
batch_dims (int, 可選) –
batch_dims(即要考慮用於batch_size的前導維度的數量)。與batch_size互斥。請注意,這是 tensordict 的 __最大__ 批次維度數,可以容忍較小的數字。names (str 列表, 可選) – 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)
- abstract from_dict_instance(input_dict, *others, auto_batch_size: Optional[bool] = None, batch_size=None, device=None, batch_dims=None, names: Optional[List[str]] = 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, optional): 讀取模式。 預設為
"r"。 auto_batch_size (bool, optional): 如果True,將自動計算批次大小。預設為
False。- batch_dims (int, optional): 如果 auto_batch_size 為
True,則定義輸出的維度數 tensordict 應該具有。 預設為
None(每個級別的完整批次大小)。
batch_size (torch.Size, optional): TensorDict 的批次大小。 預設為
None。- batch_dims (int, optional): 如果 auto_batch_size 為
- 返回:
輸入 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, 可選) – 如果為
True,將返回一個TensorDictParams例項,該例項可用於在torch.nn.Module中儲存引數。預設為False。lock (bool, 可選) – 如果為
True,則生成的 tensordict 將被鎖定。預設為True。use_state_dict (bool, 可選) –
如果為
True,則將使用模組的 state-dict,並將其展開為具有模型樹結構的 TensorDict。預設為False。注意
這在必須使用 state-dict 鉤子時特別有用。
示例
>>> from torch import nn >>> module = nn.TransformerDecoder( ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), ... num_layers=1 ... ) >>> params = TensorDict.from_module(module) >>> print(params["layers", "0", "linear1"]) TensorDict( fields={ bias: Parameter(shape=torch.Size([2048]), device=cpu, dtype=torch.float32, is_shared=False), weight: Parameter(shape=torch.Size([2048, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
- classmethod from_modules(*modules, as_module: bool = False, lock: bool = True, use_state_dict: bool = False, lazy_stack: bool = False, expand_identical: bool = False)¶
檢索多個模組的引數,用於透過 vmap 進行 ensebmle 學習/特徵應用。
- 引數:
modules (nn.Module 的序列) – 要從中獲取引數的模組。如果模組的結構不同,則需要延遲堆疊(請參閱下面的
lazy_stack引數)。- 關鍵字引數:
as_module (bool, 可選) – 如果為
True,將返回一個TensorDictParams例項,該例項可用於在torch.nn.Module中儲存引數。預設為False。lock (bool, 可選) – 如果為
True,則生成的 tensordict 將被鎖定。預設為True。use_state_dict (bool, 可選) –
如果為
True,則將使用模組的 state-dict,並將其展開為具有模型樹結構的 TensorDict。預設為False。注意
這在必須使用 state-dict 鉤子時特別有用。
lazy_stack (bool, 可選) –
引數應密集堆疊還是延遲堆疊。預設為
False(密集堆疊)。注意
lazy_stack和as_module是互斥的功能。警告
延遲輸出和非延遲輸出之間存在一個關鍵區別,即非延遲輸出將使用所需的批次大小重新例項化引數,而
lazy_stack將僅將引數表示為延遲堆疊。這意味著,雖然原始引數可以在lazy_stack=True時安全地傳遞給最佳化器,但當設定為True時,需要傳遞新的引數。警告
雖然使用延遲堆疊來保留原始引數引用可能很誘人,但請記住,每次呼叫
get()時,延遲堆疊都會執行堆疊。這將需要記憶體(引數大小的 N 倍,如果構建了圖,則更多)和計算時間。這也意味著最佳化器將包含更多引數,並且像step()或zero_grad()這樣的操作將需要更長的時間才能執行。通常,lazy_stack應保留給極少數用例。expand_identical (bool, 可選) – 如果
True並且相同的引數(相同的身份)正在堆疊到自身,則將返回此引數的擴充套件版本。當lazy_stack=True時,將忽略此引數。
示例
>>> from torch import nn >>> from tensordict import TensorDict >>> torch.manual_seed(0) >>> empty_module = nn.Linear(3, 4, device="meta") >>> n_models = 2 >>> modules = [nn.Linear(3, 4) for _ in range(n_models)] >>> params = TensorDict.from_modules(*modules) >>> print(params) TensorDict( fields={ bias: Parameter(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Parameter(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([2]), device=None, is_shared=False) >>> # example of batch execution >>> def exec_module(params, x): ... with params.to_module(empty_module): ... return empty_module(x) >>> x = torch.randn(3) >>> y = torch.vmap(exec_module, (0, None))(params, x) >>> assert y.shape == (n_models, 4) >>> # since lazy_stack = False, backprop leaves the original params untouched >>> y.sum().backward() >>> assert params["weight"].grad.norm() > 0 >>> assert modules[0].weight.grad is None
使用
lazy_stack=True時,情況略有不同>>> params = TensorDict.from_modules(*modules, lazy_stack=True) >>> print(params) LazyStackedTensorDict( fields={ bias: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, exclusive_fields={ }, batch_size=torch.Size([2]), device=None, is_shared=False, stack_dim=0) >>> # example of batch execution >>> y = torch.vmap(exec_module, (0, None))(params, x) >>> assert y.shape == (n_models, 4) >>> y.sum().backward() >>> assert modules[0].weight.grad is not None
- classmethod from_namedtuple(named_tuple, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶
遞迴地將 namedtuple 轉換為 TensorDict。
- 引數:
named_tuple – 要轉換的 namedtuple 例項。
- 關鍵字引數:
auto_batch_size (bool, 可選) – 如果
True,將自動計算批次大小。預設為False。batch_dims (int, 可選) – 如果
auto_batch_size為True,則定義輸出 tensordict 應具有的維度數。 預設為None(每個級別的完整批次大小)。device (torch.device, 可選) – 將在其上建立 TensorDict 的裝置。 預設為
None。batch_size (torch.Size, 可選) – TensorDict 的批次大小。 預設為
None。
- 返回:
輸入 namedtuple 的 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-to-tensordict 雙射變換 API。
當前接受的類包括列表、元組、命名元組和字典。
注意
對於字典,非 NestedKey 鍵將作為
NonTensorData例項單獨註冊。注意
可轉換為張量的型別(例如 int、float 或 np.ndarray)將轉換為 torch.Tensor 例項。請注意,此轉換是滿射的:將 tensordict 轉換回 pytree 將不會恢復原始型別。
示例
>>> # Create a pytree with tensor leaves, and one "weird"-looking dict key >>> class WeirdLookingClass: ... pass ... >>> weird_key = WeirdLookingClass() >>> # Make a pytree with tuple, lists, dict and namedtuple >>> pytree = ( ... [torch.randint(10, (3,)), torch.zeros(2)], ... { ... "tensor": torch.randn( ... 2, ... ), ... "td": TensorDict({"one": 1}), ... weird_key: torch.randint(10, (2,)), ... "list": [1, 2, 3], ... }, ... {"named_tuple": TensorDict({"two": torch.ones(1) * 2}).to_namedtuple()}, ... ) >>> # Build a TensorDict from that pytree >>> td = TensorDict.from_pytree(pytree) >>> # Recover the pytree >>> pytree_recon = td.to_pytree() >>> # Check that the leaves match >>> def check(v1, v2): >>> assert (v1 == v2).all() >>> >>> torch.utils._pytree.tree_map(check, pytree, pytree_recon) >>> assert weird_key in pytree_recon[1]
- classmethod from_struct_array(struct_array: ndarray, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None) T¶
將結構化的 numpy 陣列轉換為 TensorDict。
生成的 TensorDict 將與 numpy 陣列共享相同的記憶體內容(這是一個零複製操作)。就地更改結構化的 numpy 陣列的值將影響 TensorDict 的內容。
注意
此方法執行零複製操作,這意味著生成的 TensorDict 將與輸入 numpy 陣列共享相同的記憶體內容。因此,就地更改 numpy 陣列的值將影響 TensorDict 的內容。
- 引數:
struct_array (np.ndarray) – 要轉換的結構化 numpy 陣列。
- 關鍵字引數:
auto_batch_size (bool, optional) – 如果為
True,則自動計算批次大小。預設為False。batch_dims (int, 可選) – 如果
auto_batch_size為True,則定義輸出 tensordict 應具有的維度數。 預設為None(每個級別的完整批次大小)。device (torch.device, optional) –
將在其上建立 TensorDict 的裝置。預設為
None。注意
更改裝置(即,指定除
None或"cpu"之外的任何裝置)將傳輸資料,從而導致返回資料的記憶體位置發生變化。batch_size (torch.Size, optional) – TensorDict 的批次大小。預設為 None。
- 返回:
輸入結構化 numpy 陣列的 TensorDict 表示。
示例
>>> x = np.array( ... [("Rex", 9, 81.0), ("Fido", 3, 27.0)], ... dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")], ... ) >>> td = TensorDict.from_struct_array(x) >>> x_recon = td.to_struct_array() >>> assert (x_recon == x).all() >>> assert x_recon.shape == x.shape >>> # Try modifying x age field and check effect on td >>> x["age"] += 1 >>> assert (td["age"] == np.array([10, 4])).all()
- classmethod from_tuple(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶
將元組轉換為 TensorDict。
- 引數:
obj – 要轉換的元組例項。
- 關鍵字引數:
auto_batch_size (bool, optional) – 如果為
True,則自動計算批次大小。預設為False。batch_dims (int, 可選) – 如果 auto_batch_size 為
True,則定義輸出 tensordict 應具有的維度數。 預設為None(每個級別的完整批次大小)。device (torch.device, 可選) – 將在其上建立 TensorDict 的裝置。 預設為
None。batch_size (torch.Size, 可選) – TensorDict 的批次大小。 預設為
None。
- 返回:
輸入元組的 TensorDict 表示。
示例
>>> my_tuple = (1, 2, 3) >>> td = TensorDict.from_tuple(my_tuple) >>> print(td) TensorDict( fields={ 0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), 1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), 2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
- classmethod fromkeys(keys: List[NestedKey], value: Any = 0)¶
從鍵列表和一個值建立一個 tensordict。
- 引數:
keys (list of NestedKey) – 指定新字典的鍵的可迭代物件。
value (compatible type, optional) – 所有鍵的值。預設為
0。
- gather(dim: int, index: Tensor, out: Optional[T] = None) T¶
沿 dim 指定的軸收集值。
- 引數:
dim (int) – 沿其收集元素的維度
index (torch.Tensor) – 一個長整型張量,其維度數與 tensordict 的維度數匹配,但只有一個維度不同(收集維度)。其元素指的是沿所需維度收集的索引。
out (TensorDictBase, optional) – 目標 tensordict。它必須與索引具有相同的形狀。
示例
>>> td = TensorDict( ... {"a": torch.randn(3, 4, 5), ... "b": TensorDict({"c": torch.zeros(3, 4, 5)}, [3, 4, 5])}, ... [3, 4]) >>> index = torch.randint(4, (3, 2)) >>> td_gather = td.gather(dim=1, index=index) >>> print(td_gather) TensorDict( fields={ a: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 2, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 2]), device=None, is_shared=False)
Gather 保留維度名稱。
示例
>>> td.names = ["a", "b"] >>> td_gather = td.gather(dim=1, index=index) >>> td_gather.names ["a", "b"]
- gather_and_stack(dst: int, group: 'torch.distributed.ProcessGroup' | None = None) T | None¶
從各個工作程序收集 tensordict,並將它們堆疊到目標工作程序中的 self 上。
- 引數:
dst (int) – 將呼叫
gather_and_stack()的目標工作程序的 rank。group (torch.distributed.ProcessGroup, optional) – 如果設定,則指定的程序組將用於通訊。否則,將使用預設程序組。預設為
None。
示例
>>> from torch import multiprocessing as mp >>> from tensordict import TensorDict >>> import torch >>> >>> def client(): ... torch.distributed.init_process_group( ... "gloo", ... rank=1, ... world_size=2, ... init_method=f"tcp://:10003", ... ) ... # Create a single tensordict to be sent to server ... td = TensorDict( ... {("a", "b"): torch.randn(2), ... "c": torch.randn(2)}, [2] ... ) ... td.gather_and_stack(0) ... >>> def server(): ... torch.distributed.init_process_group( ... "gloo", ... rank=0, ... world_size=2, ... init_method=f"tcp://:10003", ... ) ... # Creates the destination tensordict on server. ... # The first dim must be equal to world_size-1 ... td = TensorDict( ... {("a", "b"): torch.zeros(2), ... "c": torch.zeros(2)}, [2] ... ).expand(1, 2).contiguous() ... td.gather_and_stack(0) ... assert td["a", "b"] != 0 ... print("yuppie") ... >>> if __name__ == "__main__": ... mp.set_start_method("spawn") ... ... main_worker = mp.Process(target=server) ... secondary_worker = mp.Process(target=client) ... ... main_worker.start() ... secondary_worker.start() ... ... main_worker.join() ... secondary_worker.join()
- get(key)¶
- get(key, default)
獲取與輸入鍵一起儲存的值。
- 引數:
key (str, tuple of str) – 要查詢的鍵。如果是 str 元組,則等效於連結的 getattr 呼叫。
default –
如果在 tensordict 中找不到鍵,則使用預設值。預設為
None。警告
之前,如果 tensordict 中不存在某個鍵,並且沒有傳遞預設值,則會引發 KeyError。從 v0.7 開始,此行為已更改,而是返回 None 值(與 dict.get 的行為一致)。要採用舊的行為,請設定環境變數 export TD_GET_DEFAULTS_TO_NONE=’0’ 或呼叫 :func`~tensordict.set_get_defaults_to_none(False)`。
示例
>>> td = TensorDict({"x": 1}, batch_size=[]) >>> td.get("x") tensor(1) >>> td.get("y") None
- get_at(key, index)¶
- get_at(key, index, default)
從鍵 key 的索引 idx 處獲取 tensordict 的值。
- 引數:
key (str, 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_non_tensor(key: NestedKey, default=_NoDefault.ZERO)¶
獲取非張量值(如果存在),如果未找到非張量值,則獲取 default。
此方法對張量/TensorDict 值具有魯棒性,這意味著如果收集的值是常規張量,也會返回它(儘管此方法有一些開銷,不應在其自然範圍之外使用)。
有關如何在 tensordict 中設定非張量值的更多資訊,請參閱
set_non_tensor()。- 引數:
key (NestedKey) – NonTensorData 物件的位置。
default (Any, 可選) – 如果找不到鍵,則返回的值。
- 返回:
tensordict.tensorclass.NonTensorData的內容, 或者,如果不是
tensordict.tensorclass.NonTensorData,則返回與key對應的條目(如果找不到條目,則返回default)。
示例
>>> data = TensorDict({}, batch_size=[]) >>> data.set_non_tensor(("nested", "the string"), "a string!") >>> assert data.get_non_tensor(("nested", "the string")) == "a string!" >>> # regular `get` works but returns a NonTensorData object >>> data.get(("nested", "the string")) NonTensorData( data='a string!', batch_size=torch.Size([]), device=None, is_shared=False)
- property grad¶
返回一個 tensordict,其中包含葉張量的 .grad 屬性。
- half()¶
將所有張量轉換為
torch.half。
- int()¶
將所有張量轉換為
torch.int。
- int16()¶
將所有張量轉換為
torch.int16。
- int32()¶
將所有張量轉換為
torch.int32。
- int64()¶
將所有張量轉換為
torch.int64。
- int8()¶
將所有張量轉換為
torch.int8。
- irecv(src: int, *, group: 'torch.distributed.ProcessGroup' | None = None, return_premature: bool = False, init_tag: int = 0, pseudo_rand: bool = False) tuple[int, list[torch.Future]] | list[torch.Future] | None¶
非同步接收 tensordict 的內容並使用它更新內容。
有關上下文,請檢視
isend()方法中的示例。- 引數:
src (int) – 源 worker 的 rank。
- 關鍵字引數:
group (torch.distributed.ProcessGroup, optional) – 如果設定,則指定的程序組將用於通訊。否則,將使用預設程序組。預設為
None。return_premature (bool) – 如果為
True,則返回一個 future 列表,用於等待直到 tensordict 被更新。預設為False,即等待直到在呼叫中完成更新。init_tag (int) – 源 worker 使用的
init_tag。pseudo_rand (bool) – 如果為 True,則標籤序列將是偽隨機的,允許從不同節點發送多個數據而不會重疊。請注意,生成這些偽隨機數非常昂貴(1e-5 秒/數字),這意味著它可能會減慢演算法的執行時速度。此值必須與傳遞給
isend()的值匹配。預設為False。
- 返回:
- 如果
return_premature=True,則返回一個 future 列表,用於等待 直到 tensordict 被更新。
- 如果
- is_consolidated()¶
檢查 TensorDict 是否具有合併的儲存。
- is_memmap() bool¶
檢查 tensordict 是否為記憶體對映。
如果 TensorDict 例項是記憶體對映的,則會被鎖定(條目無法重新命名、移除或新增)。如果使用全部為記憶體對映的 tensors 建立
TensorDict,這__並不__意味著is_memmap將返回True(因為新的 tensor 可能或可能不是記憶體對映的)。只有在呼叫 tensordict.memmap_() 後,該 tensordict 才會被視為記憶體對映。對於 CUDA 裝置上的 tensordicts,此值始終為
True。
檢查 tensordict 是否位於共享記憶體中。
如果 TensorDict 例項位於共享記憶體中,則會被鎖定(條目無法重新命名、移除或新增)。如果使用全部位於共享記憶體中的 tensors 建立
TensorDict,這__並不__意味著is_shared將返回True(因為新的 tensor 可能或可能不在共享記憶體中)。只有在呼叫 tensordict.share_memory_() 或將 tensordict 放置在預設情況下內容共享的裝置上(例如,"cuda")後,該 tensordict 才會被視為位於共享記憶體中。對於 CUDA 裝置上的 tensordicts,此值始終為
True。
- isend(dst: int, *, group: 'torch.distributed.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) int¶
非同步傳送 tensordict 的內容。
- 引數:
dst (int) – 應該將內容傳送到的目標 worker 的 rank。
- 關鍵字引數:
示例
>>> import torch >>> from tensordict import TensorDict >>> from torch import multiprocessing as mp >>> def client(): ... torch.distributed.init_process_group( ... "gloo", ... rank=1, ... world_size=2, ... init_method=f"tcp://:10003", ... ) ... ... td = TensorDict( ... { ... ("a", "b"): torch.randn(2), ... "c": torch.randn(2, 3), ... "_": torch.ones(2, 1, 5), ... }, ... [2], ... ) ... td.isend(0) ... >>> >>> def server(queue, return_premature=True): ... torch.distributed.init_process_group( ... "gloo", ... rank=0, ... world_size=2, ... init_method=f"tcp://:10003", ... ) ... td = TensorDict( ... { ... ("a", "b"): torch.zeros(2), ... "c": torch.zeros(2, 3), ... "_": torch.zeros(2, 1, 5), ... }, ... [2], ... ) ... out = td.irecv(1, return_premature=return_premature) ... if return_premature: ... for fut in out: ... fut.wait() ... assert (td != 0).all() ... queue.put("yuppie") ... >>> >>> if __name__ == "__main__": ... queue = mp.Queue(1) ... main_worker = mp.Process( ... target=server, ... args=(queue, ) ... ) ... secondary_worker = mp.Process(target=client) ... ... main_worker.start() ... secondary_worker.start() ... out = queue.get(timeout=10) ... assert out == "yuppie" ... main_worker.join() ... secondary_worker.join()
- isfinite() T¶
返回一個新的 tensordict,其中包含表示每個元素是否為有限值的布林元素。
當實數值不是 NaN、負無窮大或正無窮大時,它們是有限的。當複數的實部和虛部都有限時,它們是有限的。
- isnan() T¶
返回一個新的 tensordict,其中包含表示輸入的每個元素是否為 NaN 的布林元素。
當複數的實部和/或虛部為 NaN 時,則認為複數為 NaN。
- isneginf() T¶
測試輸入的每個元素是否為負無窮大。
- isposinf() T¶
測試輸入的每個元素是否為負無窮大。
- isreal() T¶
返回一個新的 tensordict,其中包含表示輸入的每個元素是否為實數值的布林元素。
- items(include_nested: bool = False, leaves_only: bool = False, is_leaf=None, *, sort: bool = False) Iterator[tuple[str, torch.Tensor]]¶
返回 tensordict 的鍵值對生成器。
- 引數:
include_nested (bool, optional) – 如果為
True,將返回巢狀值。 預設為False。leaves_only (bool, optional) – 如果為
False,則僅返回葉子節點。 預設為False。is_leaf (callable, 可選) –
一個對類型別進行呼叫的可呼叫物件,返回一個布林值,指示該類是否必須被視為葉子節點。
注意
is_leaf 的目的不是阻止對巢狀 tensordict 的遞迴呼叫,而是將某些型別標記為 “葉子”,以便在 leaves_only=True 時進行過濾。即使 is_leaf(cls) 返回 True,如果 include_nested=True,tensordict 的巢狀結構仍將被遍歷。換句話說,is_leaf 不控制遞迴深度,而是提供一種方法,當 leaves_only=True 時,從結果中過濾掉某些型別。這意味著樹中的節點既可以是葉子節點,也可以是具有子節點的節點。實際上,
is_leaf的預設值確實將 tensordict 和 tensorclass 例項從葉子集合中排除。
- 關鍵字引數:
sort (bool, optional) – 是否應對鍵進行排序。 對於巢狀鍵,鍵根據它們的連線名稱排序(即,
("a", "key")將被計數為"a.key"以進行排序)。 請注意,在處理大型 tensordicts 時,排序可能會產生顯著的開銷。 預設為False。
- abstract keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False)¶
返回 tensordict 鍵的生成器。
警告
TensorDict 的
keys()方法返回鍵的惰性檢視。 如果查詢keys但未對其進行迭代,然後修改 tensordict,則稍後迭代鍵將返回鍵的新配置。- 引數:
include_nested (bool, optional) – 如果為
True,將返回巢狀值。 預設為False。leaves_only (bool, optional) – 如果為
False,則僅返回葉子節點。 預設為False。is_leaf (callable, 可選) –
一個對類型別進行呼叫的可呼叫物件,返回一個布林值,指示該類是否必須被視為葉子節點。
注意
is_leaf 的目的不是阻止對巢狀 tensordict 的遞迴呼叫,而是將某些型別標記為 “葉子”,以便在 leaves_only=True 時進行過濾。即使 is_leaf(cls) 返回 True,如果 include_nested=True,tensordict 的巢狀結構仍將被遍歷。換句話說,is_leaf 不控制遞迴深度,而是提供一種方法,當 leaves_only=True 時,從結果中過濾掉某些型別。這意味著樹中的節點既可以是葉子節點,也可以是具有子節點的節點。實際上,
is_leaf的預設值確實將 tensordict 和 tensorclass 例項從葉子集合中排除。
- 關鍵字引數:
sort (bool, 可選) – 是否應該對鍵進行排序。對於巢狀鍵,鍵根據它們的連線名稱排序(即,
("a", "key")將被視為用於排序的"a.key")。請注意,在處理大型 tensordict 時,排序可能會產生顯著的開銷。 預設為False。
示例
>>> from tensordict import TensorDict >>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[]) >>> data.keys() ['0', '1'] >>> list(data.keys(leaves_only=True)) ['0'] >>> list(data.keys(include_nested=True, leaves_only=True)) ['0', '1', ('1', '2')]
- classmethod lazy_stack(input, dim=0, *, out=None, **kwargs)¶
建立一個 tensordict 的惰性堆疊。
有關詳細資訊,請參閱
lazy_stack()。
- lerp(end: tensordict.base.TensorDictBase | torch.Tensor, weight: tensordict.base.TensorDictBase | torch.Tensor | float)¶
基於標量或張量
weight,對兩個張量start(由self給出)和end進行線性插值。\[\text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i)\]start和end的形狀必須是可廣播的。 如果weight是一個張量,則weight、start和end的形狀必須是可廣播的。- 引數:
end (TensorDict) – 具有結束點的 tensordict。
weight (TensorDict, tensor or float) – 插值公式的權重。
- lerp_(end: tensordict.base.TensorDictBase | torch.Tensor | float, weight: tensordict.base.TensorDictBase | torch.Tensor | float)¶
lerp()的原地版本。
- lgamma() T¶
計算 TensorDict 的每個元素的
lgamma()值。
- lgamma_() T¶
原地計算 TensorDict 的每個元素的
lgamma()值。
- classmethod load(prefix: str | pathlib.Path, *args, **kwargs) T¶
從磁碟載入 tensordict。
這個類方法是
load_memmap()的代理。
- load_(prefix: str | pathlib.Path, *args, **kwargs)¶
從當前 tensordict 中的磁碟載入 tensordict。
這個類方法是
load_memmap_()的代理。
- classmethod load_memmap(prefix: str | pathlib.Path, device: Optional[device] = None, non_blocking: bool = False, *, out: Optional[TensorDictBase] = None) T¶
從磁碟載入記憶體對映的 tensordict。
- 引數:
prefix (str 或 資料夾路徑) – 儲存的 tensordict 應該從該資料夾獲取的路徑。
device (torch.device 或 等效物件, 可選) – 如果提供,資料將被非同步轉換為該裝置。支援 “meta” 裝置,在這種情況下,資料不會被載入,但會建立一組空的 “meta” 張量。這對於瞭解總模型大小和結構而無需實際開啟任何檔案非常有用。
non_blocking (bool, 可選) – 如果
True, 在將張量載入到裝置上之後,不會呼叫 synchronize。預設為False。out (TensorDictBase, 可選) – 可選的 tensordict,資料應該寫入到該 tensordict 中。
示例
>>> from tensordict import TensorDict >>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0) >>> td.memmap("./saved_td") >>> td_load = TensorDict.load_memmap("./saved_td") >>> assert (td == td_load).all()
此方法還允許載入巢狀的 tensordict。
示例
>>> nested = TensorDict.load_memmap("./saved_td/nested") >>> assert nested["e"] == 0
tensordict 也可以載入到“meta”裝置上,或者作為假張量。
示例
>>> import tempfile >>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}}) >>> with tempfile.TemporaryDirectory() as path: ... td.save(path) ... td_load = TensorDict.load_memmap(path, device="meta") ... print("meta:", td_load) ... from torch._subclasses import FakeTensorMode ... with FakeTensorMode(): ... td_load = TensorDict.load_memmap(path) ... print("fake:", td_load) meta: TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=meta, is_shared=False)}, batch_size=torch.Size([]), device=meta, is_shared=False) fake: TensorDict( fields={ a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False)}, batch_size=torch.Size([]), device=cpu, is_shared=False)
- load_memmap_(prefix: str | pathlib.Path)¶
將記憶體對映的 tensordict 的內容載入到呼叫
load_memmap_的 tensordict 中。有關更多資訊,請參閱
load_memmap()。
- load_state_dict(state_dict: OrderedDict[str, Any], strict=True, assign=False, from_flatten=False) T¶
將按照
state_dict()格式化的 state_dict 載入到 tensordict 中。- 引數:
state_dict (OrderedDict) – 要複製的 state_dict。
strict (bool, 可選) – 是否嚴格強制執行
state_dict中的鍵與此 tensordict 的torch.nn.Module.state_dict()函式返回的鍵相匹配。 預設值:Trueassign (bool, 可選) – 是否將 state dictionary 中的專案分配給 tensordict 中的相應鍵,而不是將它們就地複製到 tensordict 的當前張量中。 當
False時,保留當前模組中張量的屬性,而當True時,保留 state dict 中張量的屬性。預設值:Falsefrom_flatten (bool, 可選) – 如果
True,則假定輸入 state_dict 已被扁平化。預設為False。
示例
>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, []) >>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, []) >>> sd = data.state_dict() >>> data_zeroed.load_state_dict(sd) >>> print(data_zeroed["3", "3"]) tensor(3) >>> # with flattening >>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, []) >>> data_zeroed.load_state_dict(data.state_dict(flatten=True), from_flatten=True) >>> print(data_zeroed["3", "3"]) tensor(3)
- lock_() T¶
鎖定 tensordict 以進行非原地操作。
諸如
set()、__setitem__()、update()、rename_key_()或其他新增或刪除條目的操作將被阻止。此方法可用作裝飾器。
示例
>>> from tensordict import TensorDict >>> td = TensorDict({"a": 1, "b": 2, "c": 3}, batch_size=[]) >>> with td.lock_(): ... assert td.is_locked ... try: ... td.set("d", 0) # error! ... except RuntimeError: ... print("td is locked!") ... try: ... del td["d"] ... except RuntimeError: ... print("td is locked!") ... try: ... td.rename_key_("a", "d") ... except RuntimeError: ... print("td is locked!") ... td.set("a", 0, inplace=True) # No storage is added, moved or removed ... td.set_("a", 0) # No storage is added, moved or removed ... td.update({"a": 0}, inplace=True) # No storage is added, moved or removed ... td.update_({"a": 0}) # No storage is added, moved or removed >>> assert not td.is_locked
- log() T¶
計算 TensorDict 的每個元素的
log()值。
- log10() T¶
計算 TensorDict 的每個元素的
log10()值。
- log10_() T¶
就地計算 TensorDict 的每個元素的
log10()值。
- log1p() T¶
計算 TensorDict 的每個元素的
log1p()值。
- log1p_() T¶
就地計算 TensorDict 的每個元素的
log1p()值。
- log2() T¶
計算 TensorDict 的每個元素的
log2()值。
- log2_() T¶
計算 TensorDict 中每個元素的
log2()值,並原地修改。
- log_() T¶
計算 TensorDict 中每個元素的
log()值,並原地修改。
- logical_and(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) TensorDictBase¶
在
self和other之間執行邏輯 AND 運算。\[\text{{out}}_i = \text{{input}}_i \land \text{{other}}_i\]- 引數:
other (TensorDictBase 或 torch.Tensor) – 用於執行邏輯 AND 運算的張量或 TensorDict。
- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- logsumexp(dim=None, keepdim=False, *, out=None)¶
返回給定維度
dim中輸入 tensordict 的每一行的指數和的對數。該計算已進行數值穩定化處理。如果 keepdim 為
True,則輸出張量的大小與輸入相同,除了在維度dim中大小為1。 否則,dim被壓縮(參見squeeze()),導致輸出張量減少 1 個(或 len(dim) 個)維度。- 引數:
- 關鍵字引數:
out (TensorDictBase, 可選) – 輸出 tensordict。
- abstract make_memmap(key: NestedKey, shape: torch.Size | torch.Tensor, *, dtype: Optional[dtype] = None) MemoryMappedTensor¶
根據形狀和可能的 dtype 建立一個空的記憶體對映張量。
警告
此方法在設計上不是鎖安全的。存在於多個節點上的記憶體對映 TensorDict 例項將需要使用方法
memmap_refresh_()進行更新。寫入現有條目將導致錯誤。
- 引數:
key (NestedKey) – 要寫入的新條目的鍵。 如果 tensordict 中已存在該鍵,則會引發異常。
shape (torch.Size 或 等效, 用於巢狀張量的 torch.Tensor) – 要寫入的張量的形狀。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 新張量的 dtype。
- 返回:
一個新的記憶體對映張量。
- abstract 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 或 等效, 用於巢狀張量的 torch.Tensor) – 要寫入的張量的形狀。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 新張量的 dtype。
- 返回:
具有給定儲存的新記憶體對映張量。
- abstract make_memmap_from_tensor(key: NestedKey, tensor: Tensor, *, copy_data: bool = True) MemoryMappedTensor¶
根據給定的 tensor 建立一個空的記憶體對映 tensor。
警告
此方法在設計上不是鎖安全的。存在於多個節點上的記憶體對映 TensorDict 例項將需要使用方法
memmap_refresh_()進行更新。如果
copy_data為True,此方法總是複製儲存內容 (即,儲存不共享)。- 引數:
key (NestedKey) – 要寫入的新條目的鍵。 如果 tensordict 中已存在該鍵,則會引發異常。
tensor (torch.Tensor) – 要在物理記憶體上覆制的 tensor。
- 關鍵字引數:
copy_data (bool, 可選) – 如果為
False,則新的 tensor 將共享輸入 tensor 的元資料,例如形狀和資料型別,但內容將為空。預設為True。- 返回:
具有給定儲存的新記憶體對映張量。
- map(fn: Callable[[TensorDictBase], TensorDictBase | None], dim: int = 0, num_workers: int | None = None, *, out: TensorDictBase | 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)¶
將一個函式對映到 tensordict 在一個維度上的拆分。
此方法將透過將 tensordict 例項分成大小相等的 tensordict 並將操作分派到所需數量的工作程序上來對 tensordict 例項應用一個函式。
函式簽名應為
Callabe[[TensorDict], Union[TensorDict, Tensor]]。輸出必須支援torch.cat()操作。該函式必須是可序列化的。注意
當處理儲存在磁碟上的大型資料集(例如,記憶體對映的 tensordict)時,此方法特別有用,其中 chunk 將是原始資料的零複製切片,這些切片可以以幾乎零成本傳遞給程序。這允許以很小的成本處理非常大的資料集(例如,超過 Tb 大小的資料集)。
- 引數:
- 關鍵字引數:
out (TensorDictBase, 可選) – 用於輸出的可選容器。其沿提供的
dim維度的批次大小必須與self.ndim匹配。 如果它是共享的或記憶體對映的(is_shared()或is_memmap()返回True),它將在遠端程序中填充,避免資料向內傳輸。 否則,self切片中的資料將被髮送到程序,在當前程序上收集並原地寫入out。chunksize (int, 可選) – 每個資料塊的大小。
chunksize為 0 將沿著期望的維度解綁 tensordict,並在應用函式後重新堆疊它,而chunksize>0將分割 tensordict 並在生成的 tensordict 列表上呼叫torch.cat()。 如果未提供,則塊的數量將等於工作程序的數量。 對於非常大的 tensordict,如此大的塊可能不適合記憶體以進行操作,可能需要更多塊才能使操作實際可行。 此引數與num_chunks互斥。num_chunks (int, 可選) – 將 tensordict 分割成的塊的數量。 如果未提供,則塊的數量將等於工作程序的數量。 對於非常大的 tensordict,如此大的塊可能不適合記憶體以進行操作,可能需要更多塊才能使操作實際可行。 此引數與
chunksize互斥。pool (mp.Pool, 可選) – 用於執行作業的多程序 Pool 例項。 如果未提供,將在
map方法中建立一個池。generator (torch.Generator, 可選) –
用於播種的生成器。 將從中生成基本種子,並且池的每個工作程序將使用提供的種子(遞增一個從
0到num_workers的唯一整數)進行播種。 如果未提供生成器,則將使用一個隨機整數作為種子。 為了與未播種的工作程序一起工作,應該單獨建立一個池並直接傳遞給map()。注意
提供低值種子時應謹慎,因為這可能導致實驗之間的自相關,例如:如果請求了 8 個工作程序並且種子是 4,則工作程序種子將從 4 到 11。 如果種子是 5,則工作程序種子將從 5 到 12。 這兩個實驗將有 7 個種子的重疊,這可能會對結果產生意想不到的影響。
注意
播種工作程序的目的是讓每個工作程序都有獨立的種子,而不是在 map 方法的呼叫之間獲得可重現的結果。 換句話說,兩個實驗可能會並且很可能會返回不同的結果,因為不可能知道哪個工作程序將選擇哪個作業。 但是,我們可以確保每個工作程序都有不同的種子,並且每個工作程序上的偽隨機操作將是不相關的。
max_tasks_per_child (int, 可選) – 每個子程序選擇的最大作業數。 預設為
None,即對作業數量沒有限制。worker_threads (int, 可選) – 工作程序的執行緒數。 預設為
1。index_with_generator (bool, 可選) – 如果為
True,則 tensordict 的分割/分塊將在查詢期間完成,從而節省初始化時間。 請注意,chunk()和split()比索引(在生成器中使用)效率高得多,因此在初始化時節省處理時間可能會對總執行時間產生負面影響。 預設為False。pbar (bool, 可選) – 如果為
True,將顯示進度條。 需要提供 tqdm。 預設為False。mp_start_method (str, 可選) – 多程序的啟動方法。 如果未提供,將使用預設的啟動方法。 接受的字串是
"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[[TensorDictBase], TensorDictBase | None], dim: int = 0, num_workers: int | None = None, *, shuffle: bool = False, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: torch.Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = True, pbar: bool = False, mp_start_method: str | None = None)¶
迭代地將函式對映到 tensordict 跨一個維度的拆分。
這是
map()的可迭代版本。此方法將透過將 tensordict 例項分塊為大小相等的 tensordict 並將操作分派到所需數量的工作執行緒來將函式應用於 tensordict 例項。它將一次產生一個結果。
函式簽名應為
Callabe[[TensorDict], Union[TensorDict, Tensor]]。該函式必須是可序列化的。注意
當處理儲存在磁碟上的大型資料集(例如,記憶體對映的 tensordict)時,此方法特別有用,其中 chunk 將是原始資料的零複製切片,這些切片可以以幾乎零成本傳遞給程序。這允許以很小的成本處理非常大的資料集(例如,超過 Tb 大小的資料集)。
注意
此函式可用於表示資料集並以類似 dataloader 的方式從中載入資料。
- 引數:
- 關鍵字引數:
shuffle (bool, optional) – 指示是否應該全域性打亂索引。如果
True,則每個批次將包含非連續樣本。如果index_with_generator=False且 shuffle=True`,則會引發錯誤。預設為False。chunksize (int, 可選) – 每個資料塊的大小。
chunksize為 0 將沿著期望的維度解綁 tensordict,並在應用函式後重新堆疊它,而chunksize>0將分割 tensordict 並在生成的 tensordict 列表上呼叫torch.cat()。 如果未提供,則塊的數量將等於工作程序的數量。 對於非常大的 tensordict,如此大的塊可能不適合記憶體以進行操作,可能需要更多塊才能使操作實際可行。 此引數與num_chunks互斥。num_chunks (int, 可選) – 將 tensordict 分割成的塊的數量。 如果未提供,則塊的數量將等於工作程序的數量。 對於非常大的 tensordict,如此大的塊可能不適合記憶體以進行操作,可能需要更多塊才能使操作實際可行。 此引數與
chunksize互斥。pool (mp.Pool, 可選) – 用於執行作業的多程序 Pool 例項。 如果未提供,將在
map方法中建立一個池。generator (torch.Generator, 可選) –
用於播種的生成器。 將從中生成基本種子,並且池的每個工作程序將使用提供的種子(遞增一個從
0到num_workers的唯一整數)進行播種。 如果未提供生成器,則將使用一個隨機整數作為種子。 為了與未播種的工作程序一起工作,應該單獨建立一個池並直接傳遞給map()。注意
提供低值種子時應謹慎,因為這可能導致實驗之間的自相關,例如:如果請求了 8 個工作程序並且種子是 4,則工作程序種子將從 4 到 11。 如果種子是 5,則工作程序種子將從 5 到 12。 這兩個實驗將有 7 個種子的重疊,這可能會對結果產生意想不到的影響。
注意
播種工作程序的目的是讓每個工作程序都有獨立的種子,而不是在 map 方法的呼叫之間獲得可重現的結果。 換句話說,兩個實驗可能會並且很可能會返回不同的結果,因為不可能知道哪個工作程序將選擇哪個作業。 但是,我們可以確保每個工作程序都有不同的種子,並且每個工作程序上的偽隨機操作將是不相關的。
max_tasks_per_child (int, 可選) – 每個子程序選擇的最大作業數。 預設為
None,即對作業數量沒有限制。worker_threads (int, 可選) – 工作程序的執行緒數。 預設為
1。index_with_generator (bool, optional) –
如果
True,tensordict 的拆分/分塊將在查詢期間完成,從而節省初始化時間。 請注意,chunk()和split()比索引(在生成器中使用)效率更高,因此在初始化時節省處理時間可能會對總執行時間產生負面影響。預設為True。注意
index_with_generator的預設值對於map_iter和map不同,前者假定在記憶體中儲存 TensorDict 的拆分版本的成本過高。pbar (bool, 可選) – 如果為
True,將顯示進度條。 需要提供 tqdm。 預設為False。mp_start_method (str, 可選) – 多程序的啟動方法。 如果未提供,將使用預設的啟動方法。 接受的字串是
"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.]])
- abstract masked_fill(mask: Tensor, value: float | bool) T¶
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.]])
- abstract masked_fill_(mask: Tensor, value: float | bool) T¶
使用所需的值填充與掩碼對應的值。
- 引數:
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.]])
- abstract masked_select(mask: Tensor) T¶
遮蔽 TensorDict 的所有張量,並返回一個新的 TensorDict 例項,該例項具有指向遮蔽值的類似鍵。
- 引數:
mask (torch.Tensor) – 用於張量的布林掩碼。形狀必須與 TensorDict
batch_size匹配。
示例
>>> td = TensorDict(source={'a': torch.zeros(3, 4)}, ... batch_size=[3]) >>> mask = torch.tensor([True, False, False]) >>> td_mask = td.masked_select(mask) >>> td_mask.get("a") tensor([[0., 0., 0., 0.]])
- max(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None, return_indices: bool = True) TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的最大值。
- 引數:
- 關鍵字引數:
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.max(dim=0) max( indices=TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False), vals=TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False), batch_size=torch.Size([4]), device=None, is_shared=False) >>> td.max() TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.max(reduce=True) tensor(3.2942)
- maximum(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) T¶
計算
self和other的逐元素最大值。- 引數:
other (TensorDict 或者 Tensor) – 另一個輸入的 tensordict 或 tensor。
- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- maximum_(other: tensordict.base.TensorDictBase | torch.Tensor) T¶
maximum()的原地版本。注意
原地
maximum不支援default關鍵字引數。
- classmethod maybe_dense_stack(input, dim=0, *, out=None, **kwargs)¶
嘗試建立 tensordict 的密集堆疊,並在需要時回退到惰性堆疊。
有關詳細資訊,請參見
maybe_dense_stack()。
- mean(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的平均值。
- 引數:
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回張量的所需資料型別。 如果指定,則在執行操作之前將輸入張量轉換為 dtype。 這對於防止資料型別溢位很有用。 預設值:
None。reduce (bool, 可選) – 如果為
True,則將在所有 TensorDict 值上進行縮減,並返回單個縮減的 tensor。 預設為False。
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.mean(dim=0) TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False) >>> td.mean() TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.mean(reduce=True) tensor(-0.0547) >>> td.mean(dim="feature") TensorDict( fields={ a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False) >>> td = TensorDict( ... a=torch.ones(3, 4, 5), ... b=TensorDict( ... c=torch.ones(3, 4, 5), ... d=torch.ones(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.mean(reduce=True, dim="feature") tensor([[1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.]]) >>> td.mean(reduce=True, dim=0) tensor([[1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.]])
- memmap(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) T¶
將所有張量寫入到新的 tensordict 中對應的記憶體對映張量。
- 引數:
- 關鍵字引數:
num_threads (int, optional) – 用於寫入記憶體對映張量的執行緒數。預設為 0。
return_early (bool, optional) – 如果
True且num_threads>0,該方法將返回 tensordict 的 future。share_non_tensor (bool, optional) – 如果為
True,非張量資料將在程序之間共享,並且在單個節點內的任何 worker 上對任何 worker 進行的寫入操作(例如,原地更新或設定)將更新所有其他 worker 上的值。如果非張量葉子的數量很高(例如,共享大量的非張量資料),這可能會導致 OOM 或類似的錯誤。預設為False。existsok (bool, optional) – 如果為
False,如果張量已經存在於相同的路徑中,則會引發異常。預設為True。
然後,TensorDict 被鎖定,這意味著任何非就地 (in-place) 寫入操作都會丟擲異常(例如,重新命名、設定或刪除條目)。一旦 tensordict 被解鎖,記憶體對映屬性就會變為
False,因為跨程序標識不再得到保證。- 返回:
如果
return_early=False,則返回一個包含儲存在磁碟上的張量的新的 tensordict,否則返回一個TensorDictFuture例項。
注意
以這種方式序列化可能會在深度巢狀的 tensordict 中很慢,因此不建議在訓練迴圈內呼叫此方法。
- memmap_(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) T¶
將所有張量就地 (in-place) 寫入到對應的記憶體對映張量。
- 引數:
- 關鍵字引數:
num_threads (int, optional) – 用於寫入記憶體對映張量的執行緒數。預設為 0。
return_early (bool, 可選) – 如果
True且num_threads>0,則該方法將返回 tensordict 的 future。可以使用 future.result() 查詢生成的 tensordict。share_non_tensor (bool, optional) – 如果為
True,非張量資料將在程序之間共享,並且在單個節點內的任何 worker 上對任何 worker 進行的寫入操作(例如,原地更新或設定)將更新所有其他 worker 上的值。如果非張量葉子的數量很高(例如,共享大量的非張量資料),這可能會導致 OOM 或類似的錯誤。預設為False。existsok (bool, optional) – 如果為
False,如果張量已經存在於相同的路徑中,則會引發異常。預設為True。
然後,TensorDict 被鎖定,這意味著任何非就地 (in-place) 寫入操作都會丟擲異常(例如,重新命名、設定或刪除條目)。一旦 tensordict 被解鎖,記憶體對映屬性就會變為
False,因為跨程序標識不再得到保證。- 返回:
如果
return_early=False,則返回 self,否則返回一個TensorDictFuture例項。
注意
以這種方式序列化可能會在深度巢狀的 tensordict 中很慢,因此不建議在訓練迴圈內呼叫此方法。
- memmap_like(prefix: Optional[str] = None, copy_existing: bool = False, *, existsok: bool = True, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T¶
建立一個內容為空的、與原始 tensordict 具有相同形狀的記憶體對映 tensordict。
- 引數:
- 關鍵字引數:
num_threads (int, optional) – 用於寫入記憶體對映張量的執行緒數。預設為 0。
return_early (bool, optional) – 如果
True且num_threads>0,該方法將返回 tensordict 的 future。share_non_tensor (bool, optional) – 如果為
True,非張量資料將在程序之間共享,並且在單個節點內的任何 worker 上對任何 worker 進行的寫入操作(例如,原地更新或設定)將更新所有其他 worker 上的值。如果非張量葉子的數量很高(例如,共享大量的非張量資料),這可能會導致 OOM 或類似的錯誤。預設為False。existsok (bool, optional) – 如果為
False,如果張量已經存在於相同的路徑中,則會引發異常。預設為True。
然後,TensorDict 被鎖定,這意味著任何非就地 (in-place) 寫入操作都會丟擲異常(例如,重新命名、設定或刪除條目)。一旦 tensordict 被解鎖,記憶體對映屬性就會變為
False,因為跨程序標識不再得到保證。- 返回:
如果
return_early=False,則返回一個新的TensorDict例項,資料儲存為記憶體對映張量,否則返回一個TensorDictFuture例項。
注意
這是將一組大型緩衝區寫入磁碟的推薦方法,因為
memmap_()將複製資訊,對於大型內容來說這可能會很慢。示例
>>> td = TensorDict({ ... "a": torch.zeros((3, 64, 64), dtype=torch.uint8), ... "b": torch.zeros(1, dtype=torch.int64), ... }, batch_size=[]).expand(1_000_000) # expand does not allocate new memory >>> buffer = td.memmap_like("/path/to/dataset")
- memmap_refresh_()¶
如果記憶體對映的 tensordict 有
saved_path,則重新整理其內容。如果未與其關聯任何路徑,此方法將引發異常。
- min(dim: int | NO_DEFAULT = _NoDefault.ZERO, keepdim: bool = False, *, reduce: bool | None = None, return_indices: bool = True) TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的最小值。
- 引數:
- 關鍵字引數:
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.min(dim=0) min( indices=TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False), vals=TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False), batch_size=torch.Size([4]), device=None, is_shared=False) >>> td.min() TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.min(reduce=True) tensor(-2.9953)
- minimum(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) T¶
計算
self和other的逐元素最小值。- 引數:
other (TensorDict 或者 Tensor) – 另一個輸入的 tensordict 或 tensor。
- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- minimum_(other: tensordict.base.TensorDictBase | torch.Tensor) T¶
minimum()的原地版本。注意
原地
minimum不支援default關鍵字引數。
- mul(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) T¶
將
other乘以self。\[\text{{out}}_i = \text{{input}}_i \times \text{{other}}_i\]支援廣播、型別提升以及整數、浮點數和複數輸入。
- 引數:
other (TensorDict, Tensor or Number) – 從
self中減去的張量或數字。- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- mul_(other: tensordict.base.TensorDictBase | torch.Tensor) T¶
mul()的原地版本。注意
原地
mul不支援default關鍵字引數。
- named_apply(fn: Callable, *others: T, nested_keys: bool = False, batch_size: Optional[Sequence[int]] = None, device: torch.device | None = _NoDefault.ZERO, names: Optional[Sequence[str]] = _NoDefault.ZERO, inplace: bool = False, default: Any = _NoDefault.ZERO, filter_empty: Optional[bool] = None, propagate_lock: bool = False, call_on_nested: bool = False, out: Optional[TensorDictBase] = None, **constructor_kwargs) Optional[T]¶
將一個鍵條件可呼叫物件應用於 tensordict 中儲存的所有值,並將它們設定在一個新的 tensordict 中。
可呼叫物件的簽名必須是
Callable[Tuple[str, Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。- 引數:
fn (Callable) – 應用於 tensordict 中 (name, tensor) 對的函式。對於每個葉子節點,只會使用其葉子節點名稱(而不是完整的 NestedKey)。
*others (TensorDictBase 例項, 可選) – 如果提供,這些 tensordict 例項應具有與 self 相同的結構。
fn引數應接收與 tensordict 數量一樣多的未命名輸入,包括 self。 如果其他 tensordict 缺少條目,可以透過default關鍵字引數傳遞預設值。nested_keys (bool, optional) – 如果
True,將使用到葉子節點的完整路徑。預設為False,即僅將最後一個字串傳遞給函式。batch_size (int 序列, 可選) – 如果提供,則生成的 TensorDict 將具有所需的 batch_size。
batch_size引數應與轉換後的 batch_size 相匹配。 這是一個僅限關鍵字的引數。device (torch.device, 可選) – 結果裝置,如果有。
names (str 列表, 可選) – 新的維度名稱,以防 batch_size 被修改。
inplace (bool, optional) – 如果為 True,則進行原地修改。預設為 False。這是一個僅關鍵字引數。
default (Any, 可選) – 其他 tensordict 中缺失條目的預設值。 如果未提供,則缺失條目將引發 KeyError。
filter_empty (bool, optional) – 如果
True,則會過濾掉空的 tensordict。這也會降低計算成本,因為不會建立和銷燬空的資料結構。為了向後相容,預設為False。propagate_lock (bool, optional) – 如果
True,鎖定的 tensordict 將產生另一個鎖定的 tensordict。預設為False。call_on_nested (bool, optional) –
如果
True,該函式將被呼叫在第一層級的張量和容器(TensorDict 或 tensorclass)上。在這種情況下,func負責將其呼叫傳播到巢狀的層級。這允許在將呼叫傳播到巢狀的 tensordict 時進行細粒度的控制。如果False,該函式將只被呼叫在葉節點上,並且apply將負責將該函式分派到所有葉節點。>>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]}) >>> def mean_tensor_only(val): ... if is_tensor_collection(val): ... raise RuntimeError("Unexpected!") ... return val.mean() >>> td_mean = td.apply(mean_tensor_only) >>> def mean_any(val): ... if is_tensor_collection(val): ... # Recurse ... return val.apply(mean_any, call_on_nested=True) ... return val.mean() >>> td_mean = td.apply(mean_any, call_on_nested=True)
out (TensorDictBase, 可選) –
一個用於寫入結果的 tensordict。這可以用於避免建立新的 tensordict
>>> td = TensorDict({"a": 0}) >>> td.apply(lambda x: x+1, out=td) >>> assert (td==1).all()
警告
如果在 tensordict 上執行的操作需要訪問多個鍵才能進行單個計算,則提供一個等於
self的out引數可能會導致該操作提供靜默的錯誤結果。例如:>>> td = TensorDict({"a": 1, "b": 1}) >>> td.apply(lambda x: x+td["a"])["b"] # Right! tensor(2) >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong! tensor(3)
**constructor_kwargs – 要傳遞給 TensorDict 建構函式的其他關鍵字引數。
- 返回:
一個新的 tensordict,包含 transformed_in 張量。
示例
>>> 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)
- abstract property names¶
tensordict 的維度名稱。
可以使用
names引數在構造時設定名稱。另請參閱
refine_names(),瞭解有關在構造後如何設定名稱的詳細資訊。
- nanmean(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有非 NaN 元素的平均值。
- 引數:
dim (int, tuple of int, 可選) – 如果為
None,則返回一個無維度的 tensordict,其中包含所有葉子的平均值(如果可以計算)。如果是整數或整數元組,則僅當此維度與 tensordict 形狀相容時,才對指定的維度呼叫 mean。目前只允許使用 “feature” 字串。使用 dim=”feature” 將實現對所有特徵維度的縮減。如果 reduce=True,將返回一個形狀與 TensorDict 的批次大小相同的張量。否則,將返回一個新的 tensordict,其結構與self相同,但特徵維度已縮減。keepdim (bool) – 輸出張量是否保留維度。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回張量的所需資料型別。 如果指定,則在執行操作之前將輸入張量轉換為 dtype。 這對於防止資料型別溢位很有用。 預設值:
None。reduce (bool, 可選) – 如果為
True,則將在所有 TensorDict 值上進行縮減,並返回單個縮減的 tensor。 預設為False。
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.nanmean(dim=0) TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False) >>> td.nanmean() TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.nanmean(reduce=True) tensor(-0.0547) >>> td.nanmean(dim="feature") TensorDict( fields={ a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False) >>> td = TensorDict( ... a=torch.ones(3, 4, 5), ... b=TensorDict( ... c=torch.ones(3, 4, 5), ... d=torch.ones(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.nanmean(reduce=True, dim="feature") tensor([[1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.]]) >>> td.nanmean(reduce=True, dim=0) tensor([[1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.]])
- nansum(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有非 NaN 元素的總和。
- 引數:
dim (int, tuple of int, 可選) – 如果為
None,則返回一個無維度的 tensordict,其中包含所有葉子的總和值(如果可以計算)。如果是整數或整數元組,則僅當此維度與 tensordict 形狀相容時,才對指定的維度呼叫 sum。目前只允許使用 “feature” 字串。使用 dim=”feature” 將實現對所有特徵維度的縮減。如果 reduce=True,將返回一個形狀與 TensorDict 的批次大小相同的張量。否則,將返回一個新的 tensordict,其結構與self相同,但特徵維度已縮減。keepdim (bool) – 輸出張量是否保留維度。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回張量的所需資料型別。 如果指定,則在執行操作之前將輸入張量轉換為 dtype。 這對於防止資料型別溢位很有用。 預設值:
None。reduce (bool, 可選) – 如果為
True,則將在所有 TensorDict 值上進行縮減,並返回單個縮減的 tensor。 預設為False。
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.nansum(dim=0) TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False) >>> td.nansum() TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.nansum(reduce=True) tensor(-0.) >>> td.nansum(dim="feature") TensorDict( fields={ a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False) >>> td = TensorDict( ... a=torch.ones(3, 4, 5), ... b=TensorDict( ... c=torch.ones(3, 4, 5), ... d=torch.ones(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.nansum(reduce=True, dim="feature") tensor([[15., 15., 15., 15.], [15., 15., 15., 15.], [15., 15., 15., 15.]]) >>> td.nansum(reduce=True, dim=0) tensor([[9., 9., 9., 9., 9.], [9., 9., 9., 9., 9.], [9., 9., 9., 9., 9.], [9., 9., 9., 9., 9.]])
- property ndim: int¶
請參閱
batch_dims()。
- ndimension() int¶
請參閱
batch_dims()。
- neg() T¶
計算 TensorDict 中每個元素的
neg()值。
- neg_() T¶
就地計算 TensorDict 中每個元素的
neg()值。
- new_empty(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶
返回一個大小為
size的 TensorDict,其中包含空的張量。預設情況下,返回的 TensorDict 具有與此 tensordict 相同的
torch.dtype和torch.device。- 引數:
size (int...) – 定義輸出張量的形狀的整數列表、元組或 torch.Size。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。預設值:如果
None,則 torch.dtype 將保持不變。device (torch.device, 可選) – 返回的 tensordict 的所需裝置。預設值:如果
None,則torch.device將保持不變。requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。預設值:
False。layout (torch.layout, 可選) – 返回的 TensorDict 值的所需佈局。預設值:
torch.strided。pin_memory (bool, 可選) – 如果設定,返回的張量將在鎖頁記憶體中分配。僅適用於 CPU 張量。預設值:
False。
- new_full(size: Size, fill_value, *, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶
返回一個大小為
size且填充了 1 的 TensorDict。預設情況下,返回的 TensorDict 具有與此 tensordict 相同的
torch.dtype和torch.device。- 引數:
size (整數序列) – 定義輸出張量形狀的整數列表、元組或 torch.Size。
fill_value (標量) – 用於填充輸出張量的數字。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。預設值:如果
None,則 torch.dtype 將保持不變。device (torch.device, 可選) – 返回的 tensordict 的所需裝置。預設值:如果
None,則torch.device將保持不變。requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。預設值:
False。layout (torch.layout, 可選) – 返回的 TensorDict 值的所需佈局。預設值:
torch.strided。pin_memory (bool, 可選) – 如果設定,返回的張量將在鎖頁記憶體中分配。僅適用於 CPU 張量。預設值:
False。
- new_ones(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶
返回一個大小為
size且填充了 1 的 TensorDict。預設情況下,返回的 TensorDict 具有與此 tensordict 相同的
torch.dtype和torch.device。- 引數:
size (int...) – 定義輸出張量的形狀的整數列表、元組或 torch.Size。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。預設值:如果
None,則 torch.dtype 將保持不變。device (torch.device, 可選) – 返回的 tensordict 的所需裝置。預設值:如果
None,則torch.device將保持不變。requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。預設值:
False。layout (torch.layout, 可選) – 返回的 TensorDict 值的所需佈局。預設值:
torch.strided。pin_memory (bool, 可選) – 如果設定,返回的張量將在鎖頁記憶體中分配。僅適用於 CPU 張量。預設值:
False。
- new_tensor(data: torch.Tensor | tensordict.base.TensorDictBase, *, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, pin_memory: Optional[bool] = None)¶
返回一個以張量
data作為資料的新 TensorDict。預設情況下,返回的 TensorDict 值具有與此張量相同的
torch.dtype和torch.device。data也可以是張量集合(TensorDict或tensorclass),在這種情況下,new_tensor方法會迭代self和data的張量對。- 引數:
data (torch.Tensor 或 TensorDictBase) – 要複製的資料。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。預設值:如果
None,則 torch.dtype 將保持不變。device (torch.device, 可選) – 返回的 tensordict 的所需裝置。預設值:如果
None,則torch.device將保持不變。requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。預設值:
False。pin_memory (bool, 可選) – 如果設定,返回的張量將在鎖頁記憶體中分配。僅適用於 CPU 張量。預設值:
False。
- new_zeros(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)¶
返回一個大小為
size的 TensorDict,並用 0 填充。預設情況下,返回的 TensorDict 具有與此 tensordict 相同的
torch.dtype和torch.device。- 引數:
size (int...) – 定義輸出張量的形狀的整數列表、元組或 torch.Size。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回的 tensordict 的所需型別。預設值:如果
None,則 torch.dtype 將保持不變。device (torch.device, 可選) – 返回的 tensordict 的所需裝置。預設值:如果
None,則torch.device將保持不變。requires_grad (bool, 可選) – 如果 autograd 應該記錄返回張量的運算。預設值:
False。layout (torch.layout, 可選) – 返回的 TensorDict 值的所需佈局。預設值:
torch.strided。pin_memory (bool, 可選) – 如果設定,返回的張量將在鎖頁記憶體中分配。僅適用於 CPU 張量。預設值:
False。
- norm(*, out=None, dtype: torch.dtype | None = None)¶
計算 tensordict 中每個 tensor 的範數。
- 關鍵字引數:
out (TensorDict, 可選) – 輸出 tensordict。
dtype (torch.dtype, 可選) – 輸出資料型別 (torch>=2.4)。
- numpy()¶
將 tensordict 轉換為 numpy 陣列的(可能是巢狀的)字典。
非 tensor 資料按原樣公開。
示例
>>> from tensordict import TensorDict >>> import torch >>> data = TensorDict({"a": {"b": torch.zeros(()), "c": "a string!"}}) >>> print(data) TensorDict( fields={ a: TensorDict( fields={ b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), c: NonTensorData(data=a string!, batch_size=torch.Size([]), device=None)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> print(data.numpy()) {'a': {'b': array(0., dtype=float32), 'c': 'a string!'}}
- param_count(*, count_duplicates: bool = True) int¶
計算引數的數量(可索引項的總數),僅計算 tensor。
- 關鍵字引數:
count_duplicates (bool) – 是否將重複的張量計數為獨立的張量。 如果
False,則僅會丟棄嚴格相同的張量(來自公共基礎張量的相同檢視但不同的 id 將被計數兩次)。 預設為 True (假設每個張量都是一個單獨的副本)。
- permute(*dims: int)¶
- permute(dims: list | tuple)
返回一個 tensordict 的檢視,其批次維度根據 dims 重新排列。
- 引數:
*dims_list (int) – tensordict 的批次維度的新排序。 或者,可以提供一個整數的可迭代物件。
dims (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) T¶
在儲存的 tensor 上呼叫
pin_memory()。- 引數:
num_threads (int 或 str) – 如果提供,則用於在葉節點上呼叫
pin_memory的執行緒數。 預設為None,這會在ThreadPoolExecutor(max_workers=None)中設定大量的執行緒。 要在主執行緒上執行對pin_memory()的所有呼叫,請傳遞num_threads=0。inplace (bool, 可選) – 如果為
True,則就地修改 tensordict。 預設為False。
- pin_memory_(num_threads: int | str = 0) T¶
呼叫儲存的張量上的
pin_memory()並返回原地修改的 TensorDict。
- pop(key: NestedKey, default: Any = _NoDefault.ZERO) Tensor¶
從 tensordict 中移除並返回一個值。
如果該值不存在且未提供預設值,則會丟擲 KeyError。
- 引數:
key (str 或 巢狀鍵) – 要查詢的條目。
default (Any, 可選) – 如果找不到鍵,則返回該值。
示例
>>> td = TensorDict({"1": 1}, []) >>> one = td.pop("1") >>> assert one == 1 >>> none = td.pop("1", default=None) >>> assert none is None
- pow(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: str | torch.Tensor | None = None) T¶
計算
self中每個元素的other次方,並返回一個包含結果的張量。other可以是單個float數字、一個 Tensor 或一個TensorDict。當
other是一個張量時,input和other的形狀必須是可廣播的。- 引數:
other (float, tensor 或 tensordict) – 指數值
- 關鍵字引數:
default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- pow_(other: tensordict.base.TensorDictBase | torch.Tensor) T¶
pow()的原地版本。注意
原地
pow不支援default關鍵字引數。
- prod(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的乘積。
- 引數:
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回張量的所需資料型別。 如果指定,則在執行操作之前將輸入張量轉換為 dtype。 這對於防止資料型別溢位很有用。 預設值:
None。reduce (bool, 可選) – 如果為
True,則將在所有 TensorDict 值上進行縮減,並返回單個縮減的 tensor。 預設為False。
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.prod(dim=0) TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False) >>> td.prod() TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.prod(reduce=True) tensor(-0.) >>> td.prod(dim="feature") TensorDict( fields={ a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False) >>> td = TensorDict( ... a=torch.ones(3, 4, 5), ... b=TensorDict( ... c=torch.ones(3, 4, 5), ... d=torch.ones(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.prod(reduce=True, dim="feature") tensor([[1., 1., 1., 1.], [1., 1., 1., 1.], [1., 1., 1., 1.]]) >>> td.prod(reduce=True, dim=0) tensor([[1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.]])
- qint32()¶
將所有張量轉換為
torch.qint32。
- qint8()¶
將所有張量轉換為
torch.qint8。
- quint4x2()¶
將所有張量轉換為
torch.quint4x2。
- quint8()¶
將所有張量轉換為
torch.quint8。
- reciprocal() T¶
計算 TensorDict 中每個元素的
reciprocal()值。
- reciprocal_() T¶
原地計算 TensorDict 中每個元素的
reciprocal()值。
- record_stream(stream: Stream) T¶
將 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) – 源 worker 的 rank。
- 關鍵字引數:
- reduce(dst, op=None, async_op=False, return_premature=False, group=None) None¶
跨所有機器歸約 tensordict。
只有
rank為 dst 的程序將收到最終結果。
- refine_names(*names) T¶
根據 names 細化 self 的維度名稱。
細化是重新命名的一種特殊情況,它“提升”未命名的維度。 None 維度可以細化為具有任何名稱; 命名維度只能細化為具有相同的名稱。
由於命名張量可以與未命名張量共存,因此細化名稱提供了一種很好的方法來編寫既適用於命名張量又適用於未命名張量的命名張量感知程式碼。
names 最多可以包含一個 Ellipsis (…)。 Ellipsis 被貪婪地展開; 它被就地展開,以使用 self.names 的相應索引中的名稱將 names 填充到與 self.dim() 相同的長度。
返回:具有根據輸入命名的維度的相同 tensordict。
示例
>>> td = TensorDict({}, batch_size=[3, 4, 5, 6]) >>> tdr = td.refine_names(None, None, None, "d") >>> assert tdr.names == [None, None, None, "d"] >>> tdr = td.refine_names("a", None, None, "d") >>> assert tdr.names == ["a", None, None, "d"]
- 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")
- abstract rename_key_(old_key: NestedKey, new_key: NestedKey, safe: bool = False) T¶
使用新字串重新命名鍵,並返回具有更新鍵名的相同 tensordict。
- repeat(repeats: Size)¶
沿指定維度重複此張量。
與
expand()不同,此函式會複製張量的資料。警告
repeat()的行為與repeat()不同,但與numpy.tile()更相似。 對於類似於numpy.repeat()的運算子,請參閱repeat_interleave()。- 引數:
repeat (torch.Size, int..., int 元組 或 int 列表) – 沿每個維度重複此張量的次數。
示例
>>> import torch >>> >>> from tensordict import TensorDict >>> >>> td = TensorDict( ... { ... "a": torch.randn(3, 4, 5), ... "b": TensorDict({ ... "c": torch.randn(3, 4, 10, 1), ... "a string": "a string!", ... }, batch_size=[3, 4, 10]) ... }, batch_size=[3, 4], ... ) >>> print(td.repeat(1, 2)) TensorDict( fields={ a: Tensor(shape=torch.Size([3, 8, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ a string: NonTensorData(data=a string!, batch_size=torch.Size([3, 8, 10]), device=None), c: Tensor(shape=torch.Size([3, 8, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 8, 10]), device=None, is_shared=False)}, batch_size=torch.Size([3, 8]), device=None, is_shared=False)
- abstract repeat_interleave(repeats: torch.Tensor | int, dim: Optional[int] = None, *, output_size: Optional[int] = None) TensorDictBase¶
重複 TensorDict 的元素。
警告
這與
repeat()不同,但與numpy.repeat()類似。- 引數:
repeats (torch.Tensor 或 int) – 每個元素的重複次數。 repeats 將被廣播以適應給定軸的形狀。
dim (int, 可選) – 對值進行重複的維度。 預設情況下,使用扁平化的輸入陣列,並返回一個扁平化的輸出陣列。
- 關鍵字引數:
output_size (int, 可選) – 給定軸的總輸出大小(例如,重複次數之和)。如果給定,則避免計算 tensordict 輸出形狀所需的流同步。
- 返回:
重複的 TensorDict,其形狀與輸入相同,除了給定的軸之外。
示例
>>> import torch >>> >>> from tensordict import TensorDict >>> >>> td = TensorDict( ... { ... "a": torch.randn(3, 4, 5), ... "b": TensorDict({ ... "c": torch.randn(3, 4, 10, 1), ... "a string": "a string!", ... }, batch_size=[3, 4, 10]) ... }, batch_size=[3, 4], ... ) >>> print(td.repeat_interleave(2, dim=0)) TensorDict( fields={ a: Tensor(shape=torch.Size([6, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ a string: NonTensorData(data=a string!, batch_size=torch.Size([6, 4, 10]), device=None), c: Tensor(shape=torch.Size([6, 4, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([6, 4, 10]), device=None, is_shared=False)}, batch_size=torch.Size([6, 4]), device=None, is_shared=False)
- replace(*args, **kwargs)¶
建立一個 tensordict 的淺複製,其中的條目已被替換。
接受一個未命名引數,該引數必須是
TensorDictBase子類的字典。 此外,可以使用命名關鍵字引數更新第一級條目。- 返回:
如果輸入為非空,則返回
self的副本,其中的條目已更新。 如果提供空字典或不提供字典,並且 kwargs 為空,則返回self。
- requires_grad_(requires_grad=True) T¶
更改 autograd 是否應記錄此張量的操作:就地設定此張量的 requires_grad 屬性。
返回此 tensordict。
- 引數:
requires_grad (bool, 可選) – autograd 是否應記錄此 tensordict 上的操作。 預設為
True。
- abstract reshape(*shape: int)¶
- abstract reshape(shape: list | tuple)
返回一個連續的、重塑為所需形狀的張量。
- 引數:
*shape (int) – 結果 tensordict 的新形狀。
- 返回:
具有重塑鍵的 TensorDict
示例
>>> td = TensorDict({ ... 'x': torch.arange(12).reshape(3, 4), ... }, batch_size=[3, 4]) >>> td = td.reshape(12) >>> print(td['x']) torch.Tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
- round() T¶
計算 TensorDict 中每個元素的
round()值。
- round_() T¶
就地計算 TensorDict 中每個元素的
round()值。
- save(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T¶
將 tensordict 儲存到磁碟。
此函式是
memmap()的代理。
- property saved_path¶
返回儲存 memmap 儲存的 TensorDict 的路徑。
一旦 is_memmap() 返回
False(例如,當 tensordict 被解鎖時),此引數就會消失。
- select(*keys: NestedKey, inplace: bool = False, strict: bool = True) T¶
選擇 tensordict 的鍵,並返回一個新的 tensordict,其中只包含選定的鍵。
這些值未被複制:對原始或新 tensordict 中的張量進行原地修改會導致兩個 tensordict 都發生更改。
- 引數:
- 返回:
一個新的 tensordict (如果
inplace=True,則為同一個) ,其中只包含選定的鍵。
注意
要在 tensordict 中選擇鍵並返回此 tensordict 的一個版本,其中不包含這些鍵,請參閱
split_keys()方法。示例
>>> from tensordict import TensorDict >>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, []) >>> td.select("a", ("b", "c")) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.select("a", "b") TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.select("this key does not exist", strict=False) TensorDict( fields={ }, batch_size=torch.Size([]), device=None, is_shared=False)
- send(dst: int, *, group: 'torch.distributed.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) None¶
將 tensordict 的內容傳送到遠端工作程序。
- 引數:
dst (int) – 應該將內容傳送到的目標 worker 的 rank。
- 關鍵字引數:
示例
>>> from torch import multiprocessing as mp >>> from tensordict import TensorDict >>> import torch >>> >>> >>> def client(): ... torch.distributed.init_process_group( ... "gloo", ... rank=1, ... world_size=2, ... init_method=f"tcp://:10003", ... ) ... ... td = TensorDict( ... { ... ("a", "b"): torch.randn(2), ... "c": torch.randn(2, 3), ... "_": torch.ones(2, 1, 5), ... }, ... [2], ... ) ... td.send(0) ... >>> >>> def server(queue): ... torch.distributed.init_process_group( ... "gloo", ... rank=0, ... world_size=2, ... init_method=f"tcp://:10003", ... ) ... td = TensorDict( ... { ... ("a", "b"): torch.zeros(2), ... "c": torch.zeros(2, 3), ... "_": torch.zeros(2, 1, 5), ... }, ... [2], ... ) ... td.recv(1) ... assert (td != 0).all() ... queue.put("yuppie") ... >>> >>> if __name__=="__main__": ... queue = mp.Queue(1) ... main_worker = mp.Process(target=server, args=(queue,)) ... secondary_worker = mp.Process(target=client) ... ... main_worker.start() ... secondary_worker.start() ... out = queue.get(timeout=10) ... assert out == "yuppie" ... main_worker.join() ... secondary_worker.join()
- separates(*keys: NestedKey, default: Any = _NoDefault.ZERO, strict: bool = True, filter_empty: bool = True) T¶
從 tensordict 中原地分離指定的鍵。
參見
此方法等效於在單個拆分上呼叫
split_keys()並設定inplace=True。參見
此方法等效於呼叫
exclude(),不同之處在於它返回資料的另一個拆分。- 引數:
- 返回:
分離的 tensordict。
- 返回型別:
T
示例
>>> td = TensorDict( ... a=0, ... b=0, ... c=0, ... d=0, ... ) >>> td_a_c = td.separates("a", "c") >>> print(td_a_c) TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> print(td) TensorDict( fields={ b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
- set(key: NestedKey, item: Tensor, inplace: bool = False, *, non_blocking: bool = False, **kwargs: Any) T¶
設定新的鍵值對。
- 引數:
key (str, str 的元組) – 要設定的鍵的名稱。
item (torch.Tensor 或 等效項, TensorDictBase 例項) – 要儲存在 tensordict 中的值。
inplace (bool, optional) – 如果
True並且某個鍵與 tensordict 中已存在的鍵匹配,則該鍵值對的更新將就地發生。如果 inplace 為True並且找不到條目,則會新增它。對於更嚴格的就地操作,請改用set_()。 預設為False。
- 關鍵字引數:
non_blocking (bool, optional) – 如果
True且此複製發生在不同裝置之間,則複製可能相對於主機非同步發生。- 返回:
self
示例
>>> td = TensorDict({}, batch_size[3, 4]) >>> td.set("x", torch.randn(3, 4)) >>> y = torch.randn(3, 4, 5) >>> td.set("y", y, inplace=True) # works, even if 'y' is not present yet >>> td.set("y", torch.zeros_like(y), inplace=True) >>> assert (y==0).all() # y values are overwritten >>> td.set("y", torch.ones(5), inplace=True) # raises an exception as shapes mismatch
- set_(key: NestedKey, item: Tensor, *, non_blocking: bool = False) T¶
將值設定為現有鍵,同時保留原始儲存。
- 引數:
key (str) – 值的名稱
item (torch.Tensor 或 相容型別, TensorDictBase) – 要儲存在 tensordict 中的值
- 關鍵字引數:
non_blocking (bool, optional) – 如果
True且此複製發生在不同裝置之間,則複製可能相對於主機非同步發生。- 返回:
self
示例
>>> td = TensorDict({}, batch_size[3, 4]) >>> x = torch.randn(3, 4) >>> td.set("x", x) >>> td.set_("x", torch.zeros_like(x)) >>> assert (x == 0).all()
- set_at_(key: NestedKey, value: Tensor, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], *, non_blocking: bool = False) T¶
在
index指示的索引處就地設定值。- 引數:
key (str, str 的元組) – 要修改的鍵。
value (torch.Tensor) – 要在索引 index 處設定的值
- 關鍵字引數:
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 中註冊非 tensor 值。可以使用
TensorDictBase.get_non_tensor()或直接使用 get 來檢索該值,這將返回tensordict.tensorclass.NonTensorData物件。return: self
示例
>>> data = TensorDict({}, batch_size=[]) >>> data.set_non_tensor(("nested", "the string"), "a string!") >>> assert data.get_non_tensor(("nested", "the string")) == "a string!" >>> # regular `get` works but returns a NonTensorData object >>> data.get(("nested", "the string")) NonTensorData( data='a string!', batch_size=torch.Size([]), device=None, is_shared=False)
- setdefault(key: NestedKey, default: Tensor, inplace: bool = False) Tensor¶
如果
key不在 tensordict 中,則插入key條目,其值為default。如果
key在 tensordict 中,則返回key的值,否則返回default。- 引數:
key (str 或 巢狀鍵) – 值的名稱。
default (torch.Tensor 或 相容型別, TensorDictBase) – 如果鍵不存在,則儲存在 tensordict 中的值。
- 返回:
tensordict 中鍵的值。如果鍵先前未設定,則將為 default。
示例
>>> td = TensorDict({}, batch_size=[3, 4]) >>> val = td.setdefault("a", torch.zeros(3, 4)) >>> assert (val == 0).all() >>> val = td.setdefault("a", torch.ones(3, 4)) >>> assert (val == 0).all() # output is still 0
- property shape: Size¶
參見
batch_size。
將所有張量放置在共享記憶體中。
然後鎖定 TensorDict,這意味著任何非就地寫入操作都會引發異常(例如,重新命名、設定或刪除條目)。相反,一旦 tensordict 被解鎖,share_memory 屬性將變為
False,因為不能再保證跨程序的同一性。- 返回:
self
- sigmoid() T¶
計算 TensorDict 的每個元素的
sigmoid()值。
- sigmoid_() T¶
就地計算 TensorDict 的每個元素的
sigmoid()值。
- sign() T¶
計算 TensorDict 的每個元素的
sign()值。
- sign_() T¶
就地計算 TensorDict 的每個元素的
sign()值。
- sin() T¶
計算 TensorDict 的每個元素的
sin()值。
- sin_() T¶
就地計算 TensorDict 的每個元素的
sin()值。
- sinh() T¶
計算 TensorDict 的每個元素的
sinh()值。
- sinh_() T¶
就地計算 TensorDict 的每個元素的
sinh()值。
- size(dim: Optional[int] = None) torch.Size | int¶
返回由
dim指示的維度的大小。如果未指定
dim,則返回 TensorDict 的batch_size屬性。
- softmax(dim: int, dtype: Optional[dtype] = None)¶
將 softmax 函式應用於 tensordict 元素。
- 引數:
dim (int 或 整數元組) – 將計算 softmax 的 tensordict 維度。
dtype (torch.dtype, 可選) – 返回的張量的所需資料型別。 如果指定,則在執行操作之前將輸入張量強制轉換為 dtype。 這對於防止資料型別溢位非常有用。
- property sorted_keys: list[tensordict._nestedkey.NestedKey]¶
按字母順序返回鍵。
不支援額外的引數。
如果 TensorDict 被鎖定,則鍵將被快取,直到 tensordict 被解鎖以獲得更快的執行速度。
- abstract split(split_size: int | list[int], dim: int = 0) list[tensordict.base.TensorDictBase]¶
像 torch.split 一樣,在給定維度上以指定大小分割 TensorDict 中的每個張量。
返回一個
TensorDict例項列表,其中包含分割塊項的檢視。- 引數:
- 返回:
一個 TensorDict 的列表,其中指定了給定維度的大小。
示例
>>> td = TensorDict({ ... 'x': torch.arange(12).reshape(3, 4), ... }, batch_size=[3, 4]) >>> td0, td1 = td.split([1, 2], dim=0) >>> print(td0['x']) torch.Tensor([[0, 1, 2, 3]])
- split_keys(*key_sets, inplace=False, default: Any = _NoDefault.ZERO, strict: bool = True, reproduce_struct: bool = False) Tuple[T, ...]¶
根據一個或多個鍵集合將 tensordict 分割成子集。
該方法將返回
N+1個 tensordict,其中N是提供的引數的數量。- 引數:
key_sets (Dict 的序列[in_key, out_key] 或 鍵的列表) – 各種分割。
inplace (bool, 可選) – 如果為
True,則鍵將從self中就地刪除。預設為False。default (Any, 可選) – 缺少鍵時要返回的值。 如果未指定且
strict=True,則會引發異常。strict (bool, optional) – 如果
True,則當鍵缺失時會引發異常。 預設為True。reproduce_struct (bool, 可選) – 如果為
True,則所有返回的 tensordict 都具有與self相同的樹結構,即使某些子 tensordict 不包含任何葉子。
注意
None非張量值將被忽略且不會返回。注意
該方法不會檢查提供的列表中是否存在重複項。
示例
>>> td = TensorDict( ... a=0, ... b=0, ... c=0, ... d=0, ... ) >>> td_a, td_bc, td_d = td.split_keys(["a"], ["b", "c"]) >>> print(td_bc)
- sqrt()¶
計算
self的逐元素平方根。
- squeeze(dim: int | None = None) T¶
壓縮 -self.batch_dims+1 和 self.batch_dims-1 之間的所有維度的張量,並在新的 tensordict 中返回它們。
- 引數:
dim (Optional[int]) – 沿其壓縮的維度。如果 dim 為
None,則將壓縮所有單例維度。預設為None。
示例
>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 1, 4, 2), ... }, batch_size=[3, 1, 4]) >>> td = td.squeeze() >>> td.shape torch.Size([3, 4]) >>> td.get("x").shape torch.Size([3, 4, 2])
此操作也可以用作上下文管理器。對原始 tensordict 的更改將發生在外部,即,不會更改原始張量的內容。這也假定 tensordict 沒有被鎖定(否則,需要解鎖 tensordict)。此功能與隱式壓縮不相容。
>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 1, 4, 2), ... }, batch_size=[3, 1, 4]) >>> with td.squeeze(1) as tds: ... tds.set("y", torch.zeros(3, 4)) >>> assert td.get("y").shape == [3, 1, 4]
- classmethod stack(input, dim=0, *, out=None)¶
沿給定維度將 tensordict 堆疊到單個 tensordict 中。
此呼叫等效於呼叫
torch.stack()但與 torch.compile 相容。
- stack_from_tensordict(dim: int = 0, *, sorted: Optional[Union[bool, List[NestedKey]]] = None, out: Tensor] = None) Tensor¶
將 tensordict 的所有條目堆疊到單個張量中。
- 引數:
dim (int, 可選) – 條目應沿其堆疊的維度。
- 關鍵字引數:
sorted (bool 或 NestedKeys 的列表) – 如果為
True,則條目將按字母順序堆疊。 如果為False(預設),將使用 dict 順序。 或者,可以提供鍵名稱列表,並將相應地堆疊張量。 這會產生一些開銷,因為將根據 tensordict 中的葉名稱列表檢查鍵列表。out (torch.Tensor, 可選) – 用於堆疊操作的可選目標張量。
- stack_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) T¶
將條目堆疊到新條目中,並可能刪除原始值。
- 引數:
keys (NestedKey 的序列) – 要堆疊的條目。
- 關鍵字引數
out_key (NestedKey): 堆疊輸入的新鍵名。 keep_entries (bool, 可選): 如果為
False,則將刪除keys中的條目。預設為
False。- dim (int, 可選): 堆疊必須發生的維度。
預設為
0。
返回值: self
示例
>>> td = TensorDict(a=torch.zeros(()), b=torch.ones(())) >>> td.stack_tensors("a", "b", out_key="c") >>> assert "a" not in td >>> assert (td["c"] == torch.tensor([0, 1])).all()
- state_dict(destination=None, prefix='', keep_vars=False, flatten=False) OrderedDict[str, Any]¶
從 tensordict 生成一個 state_dict。
state_dict 的結構仍然是巢狀的,除非
flatten設定為True。一個 tensordict state_dict 包含重建 tensordict 所需的所有 tensors 和元資料(目前不支援名稱)。
- 引數:
destination (dict, 可選) – 如果提供,tensordict 的狀態將被更新到 dict 中,並返回相同的物件。否則,將建立一個
OrderedDict並返回。預設值:None。prefix (str, 可選) – 新增到 tensor 名稱的字首,用於構成 state_dict 中的鍵。預設值:
''。keep_vars (bool, 可選) – 預設情況下,state dict 中返回的
torch.Tensor項與 autograd 分離。如果設定為True,則不會執行分離。預設值:False。flatten (bool, 可選) – 是否應該使用
"."字元展平結構。預設為False。
示例
>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, []) >>> sd = data.state_dict() >>> print(sd) OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3', OrderedDict([('3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])), ('__batch_size', torch.Size([])), ('__device', None)]) >>> sd = data.state_dict(flatten=True) OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3.3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])
- std(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的標準偏差值。
- 引數:
- 關鍵字引數:
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.std(dim=0) TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False) >>> td.std() TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.std(reduce=True) tensor(1.0006) >>> td.std(dim="feature") TensorDict( fields={ a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False) >>> td = TensorDict( ... a=torch.ones(3, 4, 5), ... b=TensorDict( ... c=torch.ones(3, 4, 5), ... d=torch.ones(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.std(reduce=True, dim="feature") tensor([[0., 0., 0., 0.], [0., 0., 0., 0.], [0., 0., 0., 0.]]) >>> td.std(reduce=True, dim=0) tensor([[0., 0., 0., 0., 0.], [0., 0., 0., 0., 0.], [0., 0., 0., 0., 0.], [0., 0., 0., 0., 0.]])
- sub(other: tensordict.base.TensorDictBase | torch.Tensor | float, *, alpha: float | None = None, default: str | torch.Tensor | None = None)¶
從
self中減去other,並按alpha縮放。\[\text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i\]支援廣播、型別提升以及整數、浮點數和複數輸入。
- 引數:
other (TensorDict, Tensor or Number) – 從
self中減去的張量或數字。- 關鍵字引數:
alpha (數字) –
other的乘數。default (torch.Tensor 或 str, optional) – 用於獨佔條目的預設值。 如果未提供任何值,則兩個 tensordict 鍵列表必須完全匹配。 如果傳遞
default="intersection",則只會考慮相交的鍵集,其他鍵將被忽略。 在所有其他情況下,default將用於操作兩側的所有缺失條目。
- sub_(other: tensordict.base.TensorDictBase | torch.Tensor | float, alpha: Optional[float] = None)¶
sub()的原地 (in-place) 版本。注意
原地 (in-place)
sub不支援default關鍵字引數。
- sum(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的總和值。
- 引數:
dim (int, tuple of int, 可選) – 如果為
None,則返回一個無維度的 tensordict,其中包含所有葉子的總和值(如果可以計算)。如果是整數或整數元組,則僅當此維度與 tensordict 形狀相容時,才對指定的維度呼叫 sum。目前只允許使用 “feature” 字串。使用 dim=”feature” 將實現對所有特徵維度的縮減。如果 reduce=True,將返回一個形狀與 TensorDict 的批次大小相同的張量。否則,將返回一個新的 tensordict,其結構與self相同,但特徵維度已縮減。keepdim (bool) – 輸出張量是否保留維度。
- 關鍵字引數:
dtype (torch.dtype, 可選) – 返回張量的所需資料型別。 如果指定,則在執行操作之前將輸入張量轉換為 dtype。 這對於防止資料型別溢位很有用。 預設值:
None。reduce (bool, 可選) – 如果為
True,則將在所有 TensorDict 值上進行縮減,並返回單個縮減的 tensor。 預設為False。
示例
>>> from tensordict import TensorDict >>> import torch >>> td = TensorDict( ... a=torch.randn(3, 4, 5), ... b=TensorDict( ... c=torch.randn(3, 4, 5, 6), ... d=torch.randn(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.sum(dim=0) TensorDict( fields={ a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False) >>> td.sum() TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> td.sum(reduce=True) tensor(-0.) >>> td.sum(dim="feature") TensorDict( fields={ a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False), d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3, 4, 5]), device=None, is_shared=False)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False) >>> td = TensorDict( ... a=torch.ones(3, 4, 5), ... b=TensorDict( ... c=torch.ones(3, 4, 5), ... d=torch.ones(3, 4, 5), ... batch_size=(3, 4, 5), ... ), ... batch_size=(3, 4) ... ) >>> td.sum(reduce=True, dim="feature") tensor([[15., 15., 15., 15.], [15., 15., 15., 15.], [15., 15., 15., 15.]]) >>> td.sum(reduce=True, dim=0) tensor([[9., 9., 9., 9., 9.], [9., 9., 9., 9., 9.], [9., 9., 9., 9., 9.], [9., 9., 9., 9., 9.]])
- tan() T¶
計算 TensorDict 每個元素的
tan()值。
- tan_() T¶
原地 (in-place) 計算 TensorDict 每個元素的
tan()值。
- tanh() T¶
計算 TensorDict 每個元素的
tanh()值。
- tanh_() T¶
原地 (in-place) 計算 TensorDict 每個元素的
tanh()值。
- to(device: Optional[Union[int, device]] = ..., dtype: Optional[Union[device, str]] = ..., non_blocking: bool = ..., inplace: bool = False) T¶
- to(dtype: Union[device, str], non_blocking: bool = ...) T
- to(tensor: Tensor, non_blocking: bool = ...) T
- to(*, other: T, non_blocking: bool = ...) T
- to(*, batch_size: Size) T
將 TensorDictBase 子類對映到另一個裝置、dtype 或另一個 TensorDictBase 子類(如果允許)。
不允許將張量強制轉換為新的 dtype,因為 tensordict 不一定包含單一的張量 dtype。
- 引數:
device (torch.device, optional) – tensordict 的目標裝置。
dtype (torch.dtype, optional) – tensordict 的目標浮點型或複數 dtype。
tensor (torch.Tensor, optional) – 張量,其 dtype 和裝置是此 TensorDict 中所有張量的目標 dtype 和裝置。
- 關鍵字引數:
non_blocking (bool, optional) – 操作是否應為阻塞。
memory_format (torch.memory_format, optional) – 此 tensordict 中 4D 引數和緩衝區的目標記憶體格式。
batch_size (torch.Size, optional) – 輸出 tensordict 的結果批次大小。
other (TensorDictBase, optional) –
TensorDict 例項,其 dtype 和裝置是此 TensorDict 中所有張量的目標 dtype 和裝置。
注意
由於
TensorDictBase例項沒有 dtype,因此 dtype 從示例葉子中收集。 如果有多個 dtype,則不進行 dtype 轉換。non_blocking_pin (bool, optional) –
如果
True,則張量在傳送到裝置之前會被釘住。 這將非同步完成,但可以透過num_threads引數進行控制。注意
呼叫
tensordict.pin_memory().to("cuda")通常比tensordict.to("cuda", non_blocking_pin=True)慢得多,因為 pin_memory 在第二種情況下是非同步呼叫的。 如果張量很大且數量很多,則多執行緒pin_memory通常是有益的:當要傳送的張量太少時,生成執行緒和收集資料的開銷會超過多執行緒的優勢,如果張量很小,則迭代長列表的開銷也過大。num_threads (int or None, optional) – 如果
non_blocking_pin=True,則用於pin_memory的執行緒數。 預設情況下,將生成max(1, torch.get_num_threads())個執行緒。num_threads=0將取消 pin_memory() 呼叫的任何多執行緒處理。inplace (bool, optional) – 如果
True,資料將就地寫入相同的 tensordict。 每當構建 tensordict 受 CPU 開銷限制時,這可以顯著加快速度。 預設為False。
- 返回:
如果裝置與 tensordict 裝置不同和/或如果傳遞了 dtype,則返回一個新的 tensordict 例項。 否則,返回相同的 tensordict。
batch_size僅修改就地完成。
注意
如果 TensorDict 已合併,則生成的 TensorDict 也將合併。 每個新張量都將是強制轉換為所需裝置的合併儲存的檢視。
示例
>>> data = TensorDict({"a": 1.0}, [], device=None) >>> data_cuda = data.to("cuda:0") # casts to cuda >>> data_int = data.to(torch.int) # casts to int >>> data_cuda_int = data.to("cuda:0", torch.int) # multiple casting >>> data_cuda = data.to(torch.randn(3, device="cuda:0")) # using an example tensor >>> data_cuda = data.to(other=TensorDict({}, [], device="cuda:0")) # using a tensordict example
- to_dict(*, retain_none: bool = True) dict[str, Any]¶
返回一個字典,其鍵值對與 tensordict 的鍵值對匹配。
- 引數:
retain_none (bool) – 如果
True,則 tensorclass 例項中的None值將寫入字典中。 否則,它們將被丟棄。 預設值:True。
- to_h5(filename, **kwargs)¶
將 tensordict 轉換為具有 h5 後端的 PersistentTensorDict。
- 引數:
filename (str or path) – h5 檔案的路徑。
**kwargs – 要傳遞給
h5py.File.create_dataset()的 kwargs。
- 返回:
連結到新建立的檔案的
PersitentTensorDict例項。
示例
>>> import tempfile >>> import timeit >>> >>> from tensordict import TensorDict, MemoryMappedTensor >>> td = TensorDict({ ... "a": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000)), ... "b": {"c": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000, 3))}, ... }, [1_000_000]) >>> >>> file = tempfile.NamedTemporaryFile() >>> td_h5 = td.to_h5(file.name, compression="gzip", compression_opts=9) >>> print(td_h5) PersistentTensorDict( fields={ a: Tensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float32, is_shared=False), b: PersistentTensorDict( fields={ c: Tensor(shape=torch.Size([1000000, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([1000000]), device=None, is_shared=False)}, batch_size=torch.Size([1000000]), device=None, is_shared=False)
- to_module(module: Module, *, inplace: bool | None = None, return_swap: bool = True, swap_dest=None, use_state_dict: bool = False, non_blocking: bool = False, memo=None)¶
遞迴地將 TensorDictBase 例項的內容寫入給定的 nn.Module 屬性。
to_module也可以用作上下文管理器,以臨時使用引數/緩衝區集合填充模組(參見下面的示例)。- 引數:
module (nn.Module) – 將引數寫入的模組。
- 關鍵字引數:
inplace (bool, optional) – 如果為
True,則模組中的引數或張量將被原地更新。預設為False。return_swap (bool, optional) – 如果為
True,則返回舊的引數配置。預設為False。swap_dest (TensorDictBase, optional) – 如果
return_swap為True,則指定交換寫入的 TensorDict。use_state_dict (bool, optional) – 如果為
True,將使用 state-dict API 載入引數(包括 state-dict hooks)。預設為False。non_blocking (bool, optional) – 如果
True且此複製發生在不同裝置之間,則複製可能相對於主機非同步發生。
示例
>>> from torch import nn >>> module = nn.TransformerDecoder( ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), ... num_layers=1) >>> params = TensorDict.from_module(module) >>> params.data.zero_() >>> params.to_module(module) >>> assert (module.layers[0].linear1.weight == 0).all()
將 TensorDict 用作上下文管理器對於進行函式式呼叫非常有用: .. rubric:: 示例
>>> from tensordict import from_module >>> module = nn.TransformerDecoder( ... decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4), ... num_layers=1) >>> params = TensorDict.from_module(module) >>> params = params.data * 0 # Use TensorDictParams to remake these tensors regular nn.Parameter instances >>> with params.to_module(module): ... # Call the module with zeroed params ... y = module(*inputs) >>> # The module is repopulated with its original params >>> assert (TensorDict.from_module(module) != 0).any()
- 返回:
如果
return_swap為True,則返回包含模組值的 TensorDict;否則返回None。
- to_namedtuple(dest_cls: Optional[type] = None)¶
將 TensorDict 轉換為 namedtuple。
- 引數:
dest_cls (Type, optional) – 可選的 namedtuple 類,用於轉換。
示例
>>> from tensordict import TensorDict >>> import torch >>> data = TensorDict({ ... "a_tensor": torch.zeros((3)), ... "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3]) >>> data.to_namedtuple() GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))
- to_padded_tensor(padding=0.0, mask_key: Optional[NestedKey] = None) T¶
將所有巢狀張量轉換為填充版本,並相應地調整批處理大小。
- 引數:
padding (float) – TensorDict 中張量的填充值。預設為
0.0。mask_key (NestedKey, optional) – 如果提供,則指定寫入有效值掩碼的鍵。如果異構維度不是 TensorDict 批處理大小的一部分,則會導致錯誤。預設為
None
- to_pytree()¶
將 TensorDict 轉換為 PyTree。
如果 TensorDict 不是從 PyTree 建立的,則此方法僅返回
self,而不進行修改。有關更多資訊和示例,請參見
from_pytree()。
- to_struct_array()¶
將 TensorDict 轉換為 NumPy 結構化陣列。
在
from_struct_array()-to_struct_array()迴圈中,輸入和輸出陣列的內容應匹配。但是,to_struct_array 不會保留原始陣列的記憶體內容。參見
有關更多資訊,請參見
from_struct_array()。- 返回:
輸入 TensorDict 的 NumPy 結構化陣列表示形式。
示例
>>> import torch >>> from tensordict import TensorDict >>> td = TensorDict({'a': torch.tensor([1, 2, 3]), 'b': torch.tensor([4.0, 5.0, 6.0])}, batch_size=[3]) >>> arr = td.to_struct_array() >>> print(arr) [(1, 4.) (2, 5.) (3, 6.)]
- to_tensordict(*, retain_none: Optional[bool] = None) T¶
從 TensorDictBase 返回常規的 TensorDict 例項。
- 引數:
retain_none (bool) –
如果為
True,則 tensorclass 例項中的None值將被寫入 TensorDict。否則,它們將被丟棄。預設值:True。注意
從 v0.8 開始,預設值將切換為
False。- 返回:
包含相同值的新 TensorDict 物件。
- transpose(dim0, dim1)¶
返回一個 TensorDict,它是輸入 TensorDict 的轉置版本。給定的維度
dim0和dim1被交換。轉置 TensorDict 的原地或非原地修改也會影響原始 TensorDict,因為記憶體是共享的,並且操作會映射回原始 TensorDict。
示例
>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4]) >>> tensordict_transpose = tensordict.transpose(0, 1) >>> print(tensordict_transpose.shape) torch.Size([4, 3]) >>> tensordict_transpose.set("b",, torch.randn(4, 3)) >>> print(tensordict.get("b").shape) torch.Size([3, 4])
- trunc() T¶
計算 TensorDict 中每個元素的
trunc()值。
- trunc_() T¶
計算 TensorDict 中每個元素的
trunc()值(原地操作)。
- uint16()¶
將所有張量轉換為
torch.uint16。
- uint32()¶
將所有張量轉換為
torch.uint32。
- uint64()¶
將所有張量轉換為
torch.uint64。
- uint8()¶
將所有張量轉換為
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 維度展開為期望的形狀。
- 引數:
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) T¶
將扁平的 tensordict 遞迴地轉換為巢狀的 tensordict。
TensorDict 型別將會丟失,結果將是一個簡單的 TensorDict 例項。巢狀的 tensordict 的元資料將從根節點推斷:資料樹中的所有例項將共享相同的批次大小、維度名稱和裝置。
- 引數:
示例
>>> data = TensorDict({"a": 1, "b - c": 2, "e - f - g": 3}, batch_size=[]) >>> data.unflatten_keys(separator=" - ") TensorDict( fields={ a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), b: TensorDict( fields={ c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), e: TensorDict( fields={ f: TensorDict( fields={ g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
此方法和
unflatten_keys()在處理 state-dicts 時特別有用,因為它們可以無縫地將扁平字典轉換為模仿模型結構的資料結構。示例
>>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4)) >>> ddp_model = torch.ao.quantization.QuantWrapper(model) >>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".") >>> print(state_dict) TensorDict( fields={ module: TensorDict( fields={ 0: TensorDict( fields={ bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> model_state_dict = state_dict.get("module") >>> print(model_state_dict) TensorDict( fields={ 0: TensorDict( fields={ bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False) >>> model.load_state_dict(dict(model_state_dict.flatten_keys(".")))
- unsqueeze(dim: int) T¶
對所有張量在 -td.batch_dims 和 td.batch_dims 之間的維度進行unsqueeze操作,並將它們返回到一個新的 tensordict 中。
- 引數:
dim (int) – 進行unsqueeze操作的維度
示例
>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 4, 2), ... }, batch_size=[3, 4]) >>> td = td.unsqueeze(-2) >>> td.shape torch.Size([3, 1, 4]) >>> td.get("x").shape torch.Size([3, 1, 4, 2])
此操作也可以用作上下文管理器。 對原始 tensordict 的更改將發生在外部,即原始張量的的內容將不會被更改。 這也假定 tensordict 沒有被鎖定(否則,需要解鎖 tensordict)。
>>> td = TensorDict({ ... 'x': torch.arange(24).reshape(3, 4, 2), ... }, batch_size=[3, 4]) >>> with td.unsqueeze(-2) as tds: ... tds.set("y", torch.zeros(3, 1, 4)) >>> assert td.get("y").shape == [3, 4]
- update(input_dict_or_td: Union[dict[str, torch.Tensor], T], clone: bool = False, inplace: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None, is_leaf: Optional[Callable[[Type], bool]] = None) T¶
使用字典或另一個 TensorDict 中的值更新 TensorDict。
- 引數:
input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。
clone (bool, 可選) – 是否應在設定之前克隆輸入(張量)字典中的張量。 預設為
False。inplace (bool, 可選) – 如果
True並且鍵與 tensordict 中的現有鍵匹配,則將對該鍵值對進行原地更新。 如果找不到該條目,則將其新增。 預設為False。
- 關鍵字引數:
- 返回:
self
示例
>>> td = TensorDict({}, batch_size=[3]) >>> a = torch.randn(3) >>> b = torch.randn(3, 4) >>> other_td = TensorDict({"a": a, "b": b}, batch_size=[]) >>> td.update(other_td, inplace=True) # writes "a" and "b" even though they can't be found >>> assert td['a'] is other_td['a'] >>> other_td = other_td.clone().zero_() >>> td.update(other_td) >>> assert td['a'] is not other_td['a']
- update_(input_dict_or_td: Union[dict[str, torch.Tensor], T], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) T¶
使用字典或另一個 TensorDict 中的值原地更新 TensorDict。
與
update()不同,如果self不知道該鍵,此函式將引發錯誤。- 引數:
input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。
clone (bool, 可選) – 是否應在設定之前克隆輸入(張量)字典中的張量。 預設為
False。
- 關鍵字引數:
keys_to_update (NestedKeys 序列, 可選) – 如果提供,則僅更新
key_to_update中的鍵列表。 旨在避免呼叫data_dest.update_(data_src.select(*keys_to_update))。non_blocking (bool, optional) – 如果
True且此複製發生在不同裝置之間,則複製可能相對於主機非同步發生。
- 返回:
self
示例
>>> a = torch.randn(3) >>> b = torch.randn(3, 4) >>> td = TensorDict({"a": a, "b": b}, batch_size=[3]) >>> other_td = TensorDict({"a": a*0, "b": b*0}, batch_size=[]) >>> td.update_(other_td) >>> assert td['a'] is not other_td['a'] >>> assert (td['a'] == other_td['a']).all() >>> assert (td['a'] == 0).all()
- update_at_(input_dict_or_td: Union[dict[str, torch.Tensor], T], idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) T¶
使用來自字典或另一個 TensorDict 的值,在指定索引處就地更新 TensorDict。
與 TensorDict.update 不同,如果 TensorDict 不知道該鍵,此函式將丟擲錯誤。
- 引數:
input_dict_or_td (TensorDictBase 或 dict) – 要寫入 self 的輸入資料。
idx (int, torch.Tensor, iterable, slice) – 應該發生更新的 tensordict 的索引。
clone (bool, optional) – 是否應在設定之前克隆輸入(tensor)字典中的張量。預設為 False。
- 關鍵字引數:
keys_to_update (NestedKeys 序列, optional) – 如果提供,則僅更新
key_to_update中的鍵列表。non_blocking (bool, optional) – 如果
True且此複製發生在不同裝置之間,則複製可能相對於主機非同步發生。
- 返回:
self
示例
>>> td = TensorDict({ ... 'a': torch.zeros(3, 4, 5), ... 'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4]) >>> td.update_at_( ... TensorDict({ ... 'a': torch.ones(1, 4, 5), ... 'b': torch.ones(1, 4, 10)}, batch_size=[1, 4]), ... slice(1, 2)) TensorDict( fields={ a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32), b: Tensor(torch.Size([3, 4, 10]), dtype=torch.float32)}, batch_size=torch.Size([3, 4]), device=None, is_shared=False) >>> assert (td[1] == 1).all()
- values(include_nested: bool = False, leaves_only: bool = False, is_leaf=None, *, sort: bool = False) Iterator[Tensor]¶
返回一個生成器,表示 tensordict 的值。
- 引數:
include_nested (bool, optional) – 如果為
True,將返回巢狀值。 預設為False。leaves_only (bool, optional) – 如果為
False,則僅返回葉子節點。 預設為False。is_leaf (callable, 可選) –
一個對類型別進行呼叫的可呼叫物件,返回一個布林值,指示該類是否必須被視為葉子節點。
注意
is_leaf 的目的不是阻止對巢狀 tensordict 的遞迴呼叫,而是將某些型別標記為 “葉子”,以便在 leaves_only=True 時進行過濾。即使 is_leaf(cls) 返回 True,如果 include_nested=True,tensordict 的巢狀結構仍將被遍歷。換句話說,is_leaf 不控制遞迴深度,而是提供一種方法,當 leaves_only=True 時,從結果中過濾掉某些型別。這意味著樹中的節點既可以是葉子節點,也可以是具有子節點的節點。實際上,
is_leaf的預設值確實將 tensordict 和 tensorclass 例項從葉子集合中排除。
- 關鍵字引數:
sort (bool, optional) – 是否應對鍵進行排序。 對於巢狀鍵,鍵根據它們的連線名稱排序(即,
("a", "key")將被計數為"a.key"以進行排序)。 請注意,在處理大型 tensordicts 時,排序可能會產生顯著的開銷。 預設為False。
- var(dim: Union[int, Tuple[int], Literal['feature']] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) tensordict.base.TensorDictBase | torch.Tensor¶
返回輸入 tensordict 中所有元素的方差值。
- 引數:
dim (int, int 元組, optional) – 如果
None,則返回一個無量綱 tensordict,其中包含所有葉子的總和值(如果可以計算)。如果為整數或整數元組,則僅當此維度與 tensordict 形狀相容時,才在指定的維度上呼叫 var。目前僅允許使用 “feature” 字串。使用 dim=”feature” 將實現對所有特徵維度的縮減。如果 reduce=True,將返回一個 TensorDict 批次大小的形狀的張量。否則,將返回一個新的 tensordict,其結構與self相同,並且具有縮減的特徵維度。keepdim (bool) – 輸出張量是否保留維度。
- 關鍵字引數:
示例
>>> 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)¶
- view(dtype)
- view(shape: Size)
返回一個 tensordict,其中包含根據與 tensordict batch_size 相容的新形狀檢視的 tensors。
或者,可以提供 dtype 作為第一個未命名引數。在這種情況下,所有 tensor 將使用相應的 dtype 進行檢視。 請注意,這假設新形狀將與提供的 dtype 相容。 有關 dtype 檢視的更多資訊,請參見
view()。- 引數:
*shape (int) – 結果 tensordict 的新形狀。
dtype (torch.dtype) – 或者,用於表示 tensor 內容的 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, other, *, out=None, pad=None)¶
返回一個
TensorDict,其中包含從 self 或 other 中選擇的元素,具體取決於 condition。- 引數:
condition (BoolTensor) – 當
True(非零)時,產生self,否則產生other。other (TensorDictBase 或 Scalar) – 值(如果
other是標量)或在 condition 為False的索引處選擇的值。
- 關鍵字引數:
out (TensorDictBase, 可選) – 輸出
TensorDictBase例項。pad (標量, 可選) – 如果提供,則源或目標 tensordict 中缺少的鍵將被寫為 torch.where(mask, self, pad) 或 torch.where(mask, pad, other)。 預設為
None,即不允許缺少鍵。
- zero_() T¶
就地將 tensordict 中的所有 tensors 置零。