tensordict 包

該 TensorDict 類透過打包一個繼承自常規 pytorch 張量特性的類字典物件，簡化了在不同模組之間傳遞多個張量的過程。

TensorDictBase(*args, **kwargs)	TensorDictBase 是 TensorDict 的抽象父類，TensorDict 是一個 torch.Tensor 資料容器。
TensorDict([source, batch_size, device, ...])	一個批處理的張量字典。
LazyStackedTensorDict(*tensordicts[, ...])	一個惰性堆疊的 TensorDict。
PersistentTensorDict(*[, batch_size, ...])	持久化 TensorDict 實現。
TensorDictParams([parameters, no_convert, lock])	一個具有引數暴露功能的 TensorDictBase 包裝器。
get_defaults_to_none([set_to_none])	返回 get 預設值的狀態。

建構函式和處理器

該庫提供了幾種與 numpy 結構化陣列、命名元組或 h5 檔案等其他資料結構進行互動的方法。該庫還暴露了專門用於操作 tensordicts 的函式，如 save、load、stack 或 cat。

cat(input[, dim, out])	將tensordicts沿給定維度連線成一個tensordict。
default_is_leaf(cls)	當一個型別不是張量集合（tensordict 或 tensorclass）時，返回 True。
from_any(obj, *[, auto_batch_size, ...])	將任何物件轉換為 TensorDict。
from_consolidated(filename)	從合併檔案重建 tensordict。
from_dict(d, *[, auto_batch_size, ...])	將字典轉換為 TensorDict。
from_h5(h5_file, *[, auto_batch_size, ...])	將 HDF5 檔案轉換為 TensorDict。
from_module(module[, as_module, lock, ...])	將模組的引數和緩衝區複製到 tensordict 中。
from_modules(*modules[, as_module, lock, ...])	為 vmap 的 ensemable 學習/特徵期望應用檢索多個模組的引數。
from_namedtuple(named_tuple, *[, ...])	將命名元組轉換為 TensorDict。
from_pytree(pytree, *[, batch_size, ...])	將 pytree 轉換為 TensorDict 例項。
from_struct_array(struct_array, *[, ...])	將結構化 numpy 陣列轉換為 TensorDict。
from_tuple(obj, *[, auto_batch_size, ...])	將元組轉換為 TensorDict。
fromkeys(keys[, value])	從鍵列表和單個值建立 tensordict。
is_batchedtensor(arg0)
is_leaf_nontensor(cls)	當一個型別不是張量集合（tensordict 或 tensorclass）或是一個非張量時，返回 True。
lazy_stack(input[, dim, out])	建立 TensorDicts 的懶惰堆疊。
load(prefix[, device, non_blocking, out])	從磁碟載入 tensordict。
load_memmap(prefix[, device, non_blocking, out])	從磁碟載入記憶體對映的 tensordict。
maybe_dense_stack(input[, dim, out])	嘗試使 TensorDicts 密集堆疊，並在需要時回退到懶惰堆疊。
memmap(data[, prefix, copy_existing, ...])	將所有張量寫入記憶體對映的 Tensor 中，並放入新的 tensordict。
save(data[, prefix, copy_existing, ...])	將tensordict儲存到磁碟。
stack(input[, dim, out])	沿給定維度將 tensordicts 堆疊成一個單一的 tensordict。

TensorDict 作為上下文管理器

在需要執行某個操作然後撤銷該操作的情況下，TensorDict 可以用作上下文管理器。這包括臨時鎖定/解鎖一個 tensordict。

>>> data.lock_()  # data.set will result in an exception
>>> with data.unlock_():
...     data.set("key", value)
>>> assert data.is_locked()

或者，在包含模型引數和緩衝區的 TensorDict 例項的幫助下執行函式呼叫。

>>> params = TensorDict.from_module(module).clone()
>>> params.zero_()
>>> with params.to_module(module):
...     y = module(x)

在第一個例子中，因為我們暫時解鎖了 tensordict data，所以可以修改它。在第二個例子中，我們用 params tensordict 例項中包含的引數和緩衝區填充模型，並在該呼叫完成後重置原始引數。

記憶體對映張量

tensordict 提供了 MemoryMappedTensor 原始型別，允許您以一種方便的方式處理儲存在物理記憶體中的張量。 MemoryMappedTensor 的主要優點是其易於構建（無需處理張量的儲存）、可以處理不適合記憶體的大型連續資料、跨程序的高效（反）序列化以及對儲存張量的高效索引。

如果所有工作程序都能訪問相同的儲存（無論是在多程序還是分散式環境中），傳遞 MemoryMappedTensor 只需傳遞一個指向磁碟上檔案的引用以及一些用於重建它的額外元資料。對於索引記憶體對映張量也是如此，只要它們的儲存的資料指標與原始指標相同。

索引記憶體對映張量比從磁碟載入多個獨立檔案要快得多，並且不需要將陣列的全部內容載入到記憶體中。但是，PyTorch 張量的物理儲存不應有任何區別。

>>> my_images = MemoryMappedTensor.empty((1_000_000, 3, 480, 480), dtype=torch.unint8)
>>> mini_batch = my_images[:10]  # just reads the first 10 images of the dataset

MemoryMappedTensor(source, *[, dtype, ...])

一個記憶體對映張量。

逐點運算

Tensordict 支援各種逐點運算，允許您對其中儲存的張量執行按元素計算。這些運算類似於對常規 PyTorch 張量執行的運算。

支援的運算

當前支援以下逐點運算：

• 左加和右加（+）

• 左減和右減（-）

• 左乘和右乘（*）

• 左除和右除（/）

• 左冪（**）

許多其他運算，如 clamp()、sqrt() 等都得到支援。

執行逐點運算

您可以在兩個 Tensordicts 之間，或者在 Tensordict 與張量/標量值之間執行逐點運算。

示例 1：Tensordict-Tensordict 運算

>>> import torch
>>> from tensordict import TensorDict
>>> td1 = TensorDict(
...     a=torch.randn(3, 4),
...     b=torch.zeros(3, 4, 5),
...     c=torch.ones(3, 4, 5, 6),
...     batch_size=(3, 4),
... )
>>> td2 = TensorDict(
...     a=torch.randn(3, 4),
...     b=torch.zeros(3, 4, 5),
...     c=torch.ones(3, 4, 5, 6),
...     batch_size=(3, 4),
... )
>>> result = td1 * td2

在此示例中，\* 運算子按元素應用於 td1 和 td2 中的相應張量。

示例 2：Tensordict-Tensor 運算

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(
...     a=torch.randn(3, 4),
...     b=torch.zeros(3, 4, 5),
...     c=torch.ones(3, 4, 5, 6),
...     batch_size=(3, 4),
... )
>>> tensor = torch.randn(4)
>>> result = td * tensor

在此，\* 運算子按元素應用於 td 中的每個張量和提供的張量。該張量被廣播以匹配 Tensordict 中每個張量的形狀。

示例 3：Tensordict-Scalar 運算

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(
...     a=torch.randn(3, 4),
...     b=torch.zeros(3, 4, 5),
...     c=torch.ones(3, 4, 5, 6),
...     batch_size=(3, 4),
... )
>>> scalar = 2.0
>>> result = td * scalar

在這種情況下，\* 運算子按元素應用於 td 中的每個張量和提供的標量。

廣播規則

當在 Tensordict 與張量/標量之間執行逐點運算時，張量/標量會被廣播以匹配 Tensordict 中每個張量的形狀：張量在左側廣播以匹配 tensordict 形狀，然後單獨在右側廣播以匹配張量形狀。這遵循 PyTorch 中使用的標準廣播規則，如果您將 TensorDict 視為單個張量例項。

例如，如果您有一個形狀為 (3, 4) 的張量的 Tensordict，並將其乘以形狀為 (4,) 的張量，則該張量將在應用運算之前被廣播到形狀 (3, 4)。如果 tensordict 包含一個形狀為 (3, 4, 5) 的張量，用於乘法的張量將被廣播到右側的 (3, 4, 5) 以進行該乘法。

如果逐點運算在多個 tensordicts 之間執行，並且它們的批次大小不同，那麼它們將被廣播到一個共同的形狀。

逐點運算的效率

如果可能，將使用 torch._foreach_<op> 融合核心來加速逐點運算的計算。

處理缺失條目

當在兩個 Tensordicts 之間執行逐點運算時，它們必須具有相同的鍵。某些運算，如 add()，有一個 default 關鍵字引數，可用於與具有獨有條目的 tensordict 進行運算。如果 default=None（預設值），則兩個 Tensordicts 必須具有完全匹配的鍵集。如果 default="intersection"，則只考慮相交的鍵集，其他鍵將被忽略。在所有其他情況下，default 將用於運算雙方的所有缺失條目。

torch.autograd.grad 的梯度計算

Tensordict 擴充套件了 :pyfunc:`torch.autograd.grad`，以便可以直接在（巢狀的）TensorDict 結構上計算梯度。所有張量都被提取為元組，執行底層的 PyTorch 操作，然後結果被打包回具有相同鍵結構的新 TensorDict 中。這使得區分處理一次處理多個張量的複雜管道變得簡單明瞭。

簡單示例

>>> from tensordict import TensorDict
>>> import torch
>>> # Build input TensorDict with requires_grad
>>> inputs = TensorDict(a=torch.randn(2, 3, requires_grad=True))
>>> outputs = inputs + 1.0
>>> grads = torch.autograd.grad(outputs, inputs, torch.ones_like(outputs))
>>> print(grads)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([2, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

端到端模組

該特性在處理接受 tensordict 並返回另一個 tensordict 的函式式模組時（例如透過 tensordict.nn.TensorDictModule()）非常有用。

>>> td_in = TensorDict(a=torch.randn(2, 3, requires_grad=True))
>>> td_out = td_module(td_in)
>>> td_out = td_out.select(*td_module.out_keys)  # keep only out keys
>>> grads = torch.autograd.grad(td_out, td_in,
...                             torch.ones_like(td_out))

Utils

utils.expand_as_right(tensor, dest)	將一個張量向右擴充套件以匹配另一個張量的形狀。
utils.expand_right(tensor, shape)	將一個張量向右擴充套件以匹配所需的形狀。
utils.isin(input, reference, key[, dim])	測試 input 中 key 的每個元素在 dim 維度上是否存在於 reference 中。
utils.remove_duplicates(input, key[, dim, ...])	移除沿指定維度在 key 中重複的索引。
capture_non_tensor_stack([allow_none])	獲取捕獲非張量堆疊的當前設定。
dense_stack_tds(td_list[, dim])	密集堆疊一個 TensorDictBase 物件列表（或一個 LazyStackedTensorDict），前提是它們具有相同的結構。
is_batchedtensor(arg0)
is_tensor_collection(datatype)	檢查資料物件或型別是否是 tensordict 庫中的張量容器。
lazy_legacy([allow_none])	如果將為選定的方法使用惰性表示，則返回 True。
make_tensordict([input_dict, batch_size, ...])	從關鍵字引數或輸入字典返回一個建立的 TensorDict。
merge_tensordicts(*tensordicts[, callback_exist])	合併 tensordicts。
pad(tensordict, pad_size[, value])	在批次維度上用常量值填充 tensordict 中的所有張量，並返回一個新的 tensordict。
pad_sequence(list_of_tensordicts[, pad_dim, ...])	填充 tensordicts 列表，以便它們可以被堆疊成連續格式。
parse_tensor_dict_string(s)	將 TensorDict repr 解析為 TensorDict。
set_capture_non_tensor_stack(mode)	一個上下文管理器或裝飾器，用於控制是否應將相同的非張量資料堆疊為單個 NonTensorData 物件或 NonTensorStack。
set_lazy_legacy(mode)	設定一些方法的惰性轉換行為。
set_list_to_stack(mode)	上下文管理器和裝飾器，用於控制 TensorDict 中列表處理的行為。
list_to_stack([allow_none])	檢索 TensorDict 中列表到堆疊轉換的當前設定。