NonTensorData¶

class tensordict.NonTensorData(data: 'Any', _metadata: 'dict | None' = None, _is_non_tensor: 'bool' = True, *, batch_size, device=None, names=None)¶

property batch_size: Size¶

检索张量类的批大小。

返回:: 批大小 (torch.Size)

property device: device¶: 检索张量类的设备类型。

classmethod from_tensordict(tensordict, non_tensordict=None)¶

张量类包装器，用于实例化新的张量类对象。

参数:

tensordict (TensorDict) – 张量类型的字典
non_tensordict (dict) – 包含非张量和嵌套张量类对象的字典

get(key: NestedKey, default: Any = _NoDefault.ZERO)¶

获取使用输入键存储的值。

参数:

key (str, str 的元组) – 要查询的键。如果为 str 的元组，则等效于 getattr 的链式调用。
default – 如果在 tensorclass 中找不到键，则为默认值。

返回:

使用输入键存储的值

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

从磁盘加载 tensordict。

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

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

在当前 tensordict 中从磁盘加载 tensordict。

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

classmethod load_memmap(prefix: str | Path, device: torch.device | None = None, non_blocking: bool = False, *, out: TensorDictBase | None = None) → T¶

从磁盘加载内存映射的 tensordict。

参数:

prefix (str 或 文件夹路径) – 保存的 tensordict 应从中获取的文件夹的路径。
device (torch.device 或 等效项, 可选) – 如果提供，数据将异步转换为该设备。支持 “meta” 设备，在这种情况下，数据不会加载，而是创建一组空的“meta”张量。这对于在不实际打开任何文件的情况下了解模型的总大小和结构很有用。
non_blocking (bool, 可选) – 如果为True，则在设备上加载张量后不会调用同步操作。默认为False。
out (TensorDictBase, 可选) – 可选的tensordict，数据将写入其中。

示例

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

此方法还允许加载嵌套的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_state_dict(state_dict: dict[str, Any], strict=True, assign=False, from_flatten=False)¶: 尝试在目标张量类上就地加载state_dict。

maybe_to_stack()¶: 如果NonTensorData对象具有非空的批大小，则将其转换为NonTensorStack对象。

memmap(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

将所有张量写入新tensordict中对应的内存映射张量。

参数:

prefix (str) – 存储内存映射张量的目录前缀。目录树结构将模仿tensordict的结构。
copy_existing (bool) – 如果为False（默认），则如果tensordict中的条目已经是磁盘上存储的张量，并且具有关联的文件，但未根据前缀保存在正确的位置，则会引发异常。如果为True，则任何现有的张量都将复制到新位置。

关键字参数：

num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为0。
return_early (bool, 可选) – 如果为True且num_threads>0，则该方法将返回tensordict的future。
share_non_tensor (bool, 可选) – 如果为True，则非张量数据将在进程之间共享，并且在单个节点内的任何工作程序上的写入操作（例如就地更新或设置）将更新所有其他工作程序上的值。如果非张量叶节点的数量很高（例如，共享大型的非张量数据堆栈），这可能会导致内存不足或类似的错误。默认为False。

然后锁定TensorDict，这意味着任何非就地写入操作都将引发异常（例如，重命名、设置或删除条目）。解锁tensordict后，内存映射属性将变为False，因为跨进程身份不再保证。

返回:: 如果return_early=False，则为一个包含磁盘上存储的张量的新的tensordict，否则为TensorDictFuture实例。

注意

以这种方式序列化对于深度嵌套的tensordict来说可能很慢，因此不建议在训练循环中调用此方法。

memmap_(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

就地将所有张量写入对应的内存映射张量。

参数:

prefix (str) – 存储内存映射张量的目录前缀。目录树结构将模仿tensordict的结构。
copy_existing (bool) – 如果为False（默认），则如果tensordict中的条目已经是磁盘上存储的张量，并且具有关联的文件，但未根据前缀保存在正确的位置，则会引发异常。如果为True，则任何现有的张量都将复制到新位置。

关键字参数：

num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为0。
return_early (bool, 可选) – 如果为True且num_threads>0，则该方法将返回tensordict的future。可以使用future.result()查询结果tensordict。
share_non_tensor (bool, 可选) – 如果为True，则非张量数据将在进程之间共享，并且在单个节点内的任何工作程序上的写入操作（例如就地更新或设置）将更新所有其他工作程序上的值。如果非张量叶节点的数量很高（例如，共享大型的非张量数据堆栈），这可能会导致内存不足或类似的错误。默认为False。

然后锁定TensorDict，这意味着任何非就地写入操作都将引发异常（例如，重命名、设置或删除条目）。解锁tensordict后，内存映射属性将变为False，因为跨进程身份不再保证。

返回:: 如果return_early=False，则为self，否则为TensorDictFuture实例。

注意

以这种方式序列化对于深度嵌套的tensordict来说可能很慢，因此不建议在训练循环中调用此方法。

memmap_like(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T¶

创建一个与原始tensordict具有相同形状的无内容内存映射tensordict。

参数:

prefix (str) – 存储内存映射张量的目录前缀。目录树结构将模仿tensordict的结构。
copy_existing (bool) – 如果为False（默认），则如果tensordict中的条目已经是磁盘上存储的张量，并且具有关联的文件，但未根据前缀保存在正确的位置，则会引发异常。如果为True，则任何现有的张量都将复制到新位置。

关键字参数：

num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为0。
return_early (bool, 可选) – 如果为True且num_threads>0，则该方法将返回tensordict的future。
share_non_tensor (bool, 可选) – 如果为True，则非张量数据将在进程之间共享，并且在单个节点内的任何工作程序上的写入操作（例如就地更新或设置）将更新所有其他工作程序上的值。如果非张量叶节点的数量很高（例如，共享大型的非张量数据堆栈），这可能会导致内存不足或类似的错误。默认为False。

然后锁定TensorDict，这意味着任何非就地写入操作都将引发异常（例如，重命名、设置或删除条目）。解锁tensordict后，内存映射属性将变为False，因为跨进程身份不再保证。

返回:: 如果 return_early=False，则返回一个新的 TensorDict 实例，其数据存储为内存映射张量；否则返回一个 TensorDictFuture 实例。

注意

这是将一组大型缓冲区写入磁盘的推荐方法，因为 memmap_() 会复制信息，对于大型内容来说，这可能很慢。

示例

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

memmap_refresh_()¶

如果内存映射的 TensorDict 具有 saved_path，则刷新其内容。

如果未关联任何路径，则此方法将引发异常。

property names: Size¶

检索张量类的维度名称。

返回:: names (字符串列表)

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

将 TensorDict 保存到磁盘。

此函数是 memmap() 的代理。

set(key: NestedKey, value: Any, inplace: bool = False, non_blocking: bool = False)¶

设置新的键值对。

参数:

key (str, str 元组) – 要设置的键的名称。如果为 str 元组，则等效于 getattr 的链式调用，最后跟一个 setattr。
value (Any) – 要存储在 TensorClass 中的值
inplace (bool, 可选) – 如果为 True，则 set 将尝试就地更新值。如果为 False 或键不存在，则值将简单地写入其目标位置。

返回:

self

state_dict(destination=None, prefix='', keep_vars=False, flatten=False) → dict[str, Any]¶: 返回一个 state_dict 字典，可用于保存和加载 TensorClass 的数据。

tolist()¶

如果批大小非空，则将数据转换为列表。

如果批大小为空，则返回数据。

unbind(dim: int)¶

返回沿指定维度解绑的 TensorClass 实例的元组。

生成的 TensorClass 实例将共享初始 TensorClass 实例的存储。

NonTensorData¶

文档

教程

资源