tensordict 包

TensorDict 类通过将多个张量打包到类似字典的对象中来简化从模块到模块传递多个张量的过程，该对象继承了常规 pytorch 张量的特性。

TensorDictBase()	TensorDictBase 是 TensorDicts 的抽象父类，TensorDicts 是 torch.Tensor 数据容器。
TensorDict([source, batch_size, device, ...])	张量的批量字典。
LazyStackedTensorDict(*tensordicts[, ...])	TensorDicts 的惰性堆叠。
PersistentTensorDict(*[, batch_size, ...])	持久 TensorDict 实现。
TensorDictParams(parameters, *[, ...])	保存包含参数的 TensorDictBase 实例。
get_defaults_to_none([set_to_none])	返回 get 默认值的状态。

构造函数和处理程序

该库提供了一些方法来与其他数据结构（如 numpy 结构化数组、namedtuples 或 h5 文件）进行交互。该库还公开了专用函数来操作 tensordict，例如 save、load、stack 或 cat。

cat(input[, dim, out])	沿给定维度将 tensordict 连接到单个 tensordict 中。
from_consolidated(filename)	从合并文件中重建 tensordict。
from_dict(input_dict[, batch_size, device, ...])	返回从字典或另一个 TensorDict 创建的 TensorDict。
from_h5(filename[, mode])	从 h5 文件创建 PersistentTensorDict。
from_module(module[, as_module, lock, ...])	将模块的参数和缓冲区复制到 tensordict 中。
from_modules(*modules[, as_module, lock, ...])	检索多个模块的参数，用于通过 vmap 进行 ensebmle 学习/期望应用程序。
from_namedtuple(named_tuple, *[, ...])	将 namedtuple 递归转换为 TensorDict。
from_pytree(pytree, *[, batch_size, ...])	将 pytree 转换为 TensorDict 实例。
from_struct_array(struct_array[, device])	将结构化的 numpy 数组转换为 TensorDict。
fromkeys(keys[, value])	从键列表和单个值创建 tensordict。
is_batchedtensor(arg0)
lazy_stack(input[, dim, out])	创建 tensordict 的惰性堆叠。
load(prefix[, device, non_blocking, out])	从磁盘加载 tensordict。
load_memmap(prefix[, device, non_blocking, out])	从磁盘加载内存映射的 tensordict。
maybe_dense_stack(input[, dim, out])	尝试对 tensordict 进行密集堆叠，并在需要时回退到惰性堆叠。
memmap(data[, prefix, copy_existing, ...])	将所有张量写入新 tensordict 中的相应内存映射张量。
save(data[, prefix, copy_existing, ...])	将 tensordict 保存到磁盘。
stack(input[, dim, out])	沿给定维度将 tensordict 堆叠到单个 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, ...])

内存映射张量。

实用工具

utils.expand_as_right(tensor, dest)	在右侧扩展张量以匹配另一个张量形状。
utils.expand_right(tensor, shape)	在右侧扩展张量以匹配所需的形状。
utils.isin(input, reference, key[, dim])	测试输入 dim 中 key 的每个元素是否也存在于参考中。
utils.remove_duplicates(input, key[, dim, ...])	删除指定维度中 key 中重复的索引。
is_batchedtensor(arg0)
is_tensor_collection(datatype)	检查数据对象或类型是否为 tensordict 库中的张量容器。
make_tensordict([input_dict, batch_size, device])	返回从关键字参数或输入字典创建的 TensorDict。
merge_tensordicts(*tensordicts[, callback_exist])	将 tensordict 合并在一起。
pad(tensordict, pad_size[, value])	使用常数值沿批次维度填充 tensordict 中的所有张量，并返回新的 tensordict。
pad_sequence(list_of_tensordicts[, pad_dim, ...])	填充 tensordict 列表，以便将它们以连续格式堆叠在一起。
dense_stack_tds(td_list[, dim])	密集堆叠 TensorDictBase 对象列表（或 LazyStackedTensorDict），前提是它们具有相同的结构。
set_lazy_legacy(mode)	将某些方法的行为设置为惰性转换。
lazy_legacy([allow_none])	如果惰性表示将用于选定的方法，则返回 True。
parse_tensor_dict_string(s)	将 TensorDict repr 解析为 TensorDict。