保存 TensorDict 和 tensorclass 对象¶

虽然我们可以使用 save() 来保存 tensordict，但这会创建一个包含数据结构整个内容的文件。很容易想象某些情况这种做法并不是最优的！

TensorDict 的序列化 API 主要依赖于 MemoryMappedTensor，它被用于将张量独立地写入磁盘，并创建一个模仿 TensorDict 的数据结构。

TensorDict 的序列化速度比 PyTorch 通过 save() 依赖 pickle 的方式快一个数量级。本文档将解释如何使用 TensorDict 创建和交互存储在磁盘上的数据。

保存内存映射的 TensorDict¶

当 tensordict 被转储为 mmap 数据结构时，每个条目对应一个单独的 *.memmap 文件，目录结构由键结构决定：通常，嵌套的键对应子目录。

将数据结构保存为结构化的内存映射张量集具有以下优点：

保存的数据可以部分加载。如果一个大型模型保存在磁盘上，但只有一个部分权重需要加载到一个单独脚本中创建的模块中，那么只有这些权重会被加载到内存中。
保存数据是安全的：使用 pickle 库序列化大型数据结构可能是不安全的，因为反序列化可以执行任意代码。TensorDict 的加载 API 只读取从保存的 json 文件和磁盘上保存的内存缓冲区中预先选择的字段。
保存速度快：由于数据写入到多个独立文件中，我们可以通过启动多个并发线程来分摊 IO 开销，每个线程各自访问一个专用文件。
保存数据的结构是显而易见的：目录树指示了数据的内容。

然而，这种方法也有一些缺点：

并非所有数据类型都可以保存。tensorclass 允许保存任何非张量数据：如果这些数据可以表示在 json 文件中，则会使用 json 格式。否则，非张量数据将作为后备方案独立保存，使用 save()。 NonTensorData 类可用于在常规 TensorDict 实例中表示非张量数据。

tensordict 的内存映射 API 依赖于四个核心方法：memmap_()、memmap()、memmap_like() 和 load_memmap()。

方法 memmap_() 和 memmap() 会在不修改或修改包含数据的 tensordict 实例的情况下将数据写入磁盘。这些方法可用于序列化模型到磁盘（我们使用多线程加速序列化）。

>>> model = nn.Transformer()
>>> weights = TensorDict.from_module(model)
>>> weights_disk = weights.memmap("/path/to/saved/dir", num_threads=32)
>>> new_weights = TensorDict.load_memmap("/path/to/saved/dir")
>>> assert (weights_disk == new_weights).all()

方法 memmap_like() 用于需要在磁盘上预先分配数据集的情况，典型用法如下：

>>> def make_datum(): # used for illustration purposes
...    return TensorDict({"image": torch.randint(255, (3, 64, 64)), "label": 0}, batch_size=[])
>>> dataset_size = 1_000_000
>>> datum = make_datum() # creates a single instance of a TensorDict datapoint
>>> data = datum.expand(dataset_size) # does NOT require more memory usage than datum, since it's only a view on datum!
>>> data_disk = data.memmap_like("/path/to/data")  # creates the two memory-mapped tensors on disk
>>> del data # data is not needed anymore

如上所示，在将 `TensorDict` 的条目转换为 `MemoryMappedTensor` 时，可以控制内存映射在磁盘上的保存位置，以便它们能够持久化并在以后加载。另一方面，文件系统也可以使用。要使用此功能，只需在上述三个序列化方法中省略 `prefix` 参数。

当指定了 `prefix` 时，数据结构将遵循 TensorDict 的结构：

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.rand(10), "b": {"c": torch.rand(10)}}, [10])
>>> td.memmap_(prefix="tensordict")

产生以下目录结构：

tensordict
├── a.memmap
├── b
│   ├── c.memmap
│   └── meta.json
└── meta.json

`meta.json` 文件包含重建 tensordict 的所有相关信息，例如设备、批次大小，还包括 tensordict 的子类型。这意味着 `TensorDict.load_memmap()` 将能够重建复杂的嵌套结构，其中子 tensordict 的类型与父级不同。

>>> from tensordict import TensorDict, tensorclass, TensorDictBase
>>> from tensordict.utils import print_directory_tree
>>> import torch
>>> import tempfile
>>> td_list = [TensorDict({"item": i}, batch_size=[]) for i in range(4)]
>>> @tensorclass
... class MyClass:
...     data: torch.Tensor
...     metadata: str
>>> tc = MyClass(torch.randn(3), metadata="some text", batch_size=[])
>>> data = TensorDict({"td_list": torch.stack(td_list), "tensorclass": tc}, [])
>>> with tempfile.TemporaryDirectory() as tempdir:
...     data.memmap_(tempdir)
...
...     loaded_data = TensorDictBase.load_memmap(tempdir)
...     assert (loaded_data == data).all()
...     print_directory_tree(tempdir)
tmpzy1jcaoq/
    tensorclass/
        _tensordict/
            data.memmap
            meta.json
        meta.json
    td_list/
        0/
            item.memmap
            meta.json
        1/
            item.memmap
            meta.json
        3/
            item.memmap
            meta.json
        2/
            item.memmap
            meta.json
        meta.json
    meta.json

处理现有的 `MemoryMappedTensor`¶

如果 `TensorDict` 已包含 `MemoryMappedTensor` 条目，则有几种可能的行为。

如果未指定 `prefix` 且多次调用 `memmap()`，则生成的 `TensorDict` 将包含与原始数据相同的数据。

>>> td = TensorDict({"a": 1}, [])
>>> td0 = td.memmap()
>>> td1 = td0.memmap()
>>> td0["a"] is td1["a"]
True

如果指定了 `prefix` 且该 `prefix` 与现有 `MemoryMappedTensor` 实例的 `prefix` 不同，则会引发异常，除非传递了 `copy_existing=True`。

>>> with tempfile.TemporaryDirectory() as tmpdir_0:
...     td0 = td.memmap(tmpdir_0)
...     td0 = td.memmap(tmpdir_0)  # works, results are just overwritten
...     with tempfile.TemporaryDirectory() as tmpdir_1:
...         td1 = td0.memmap(tmpdir_1)
...         td_load = TensorDict.load_memmap(tmpdir_1)  # works!
...     assert (td_load == td).all()
...     with tempfile.TemporaryDirectory() as tmpdir_1:
...         td_load = TensorDict.load_memmap(tmpdir_1)  # breaks!

此功能是为了防止用户无意中将内存映射张量从一个位置复制到另一个位置而实现的。

TorchSnapshot 兼容性¶

警告

由于 torchsnapshot 的维护已停止。因此，我们不会为 tensordict 与此库的兼容性实现新功能。

TensorDict 与 PyTorch 检查点库 torchsnapshot 兼容。TorchSnapshot 将独立保存您的每个张量，其数据结构模仿您的 tensordict 或 tensorclass。此外，TensorDict 自然内置了用于在磁盘上保存和加载大型数据集而无需将完整张量加载到内存中的工具：换句话说，tensordict + torchsnapshot 的组合使得有可能将一个大小为数百 GB 的张量加载到一个预先分配的 `MemmapTensor` 中，而无需一次性将其全部加载到 RAM 中。

主要有两种用例：保存和加载适合内存的 tensordict，以及保存和加载使用 `MemmapTensor` 存储在磁盘上的 tensordict。

通用用例：内存加载¶

如果目标 tensordict 未预先分配，则此方法适用。它提供了灵活性（您可以将任何 tensordict 加载到您的 tensordict 中，无需提前知道其内容），并且此方法的编码比其他方法略微容易。但是，如果您的张量非常大且不适合放入内存，这可能会失败。此外，它不允许您直接加载到您选择的设备。

用于保存操作的两个主要命令是：

>>> state = {"state": tensordict_source}
>>> snapshot = torchsnapshot.Snapshot.take(app_state=state, path="/path/to/my/snapshot")

要加载到目标 tensordict，您可以简单地加载快照并更新 tensordict。在底层，此方法将调用 `tensordict_target.load_state_dict(state_dict)`，这意味着 `state_dict` 将首先完全放入内存，然后加载到目标 tensordict 中。

>>> snapshot = Snapshot(path="/path/to/my/snapshot")
>>> state_target = {"state": tensordict_target}
>>> snapshot.restore(app_state=state_target)

这是一个完整的示例：

>>> import uuid
>>> import torchsnapshot
>>> from tensordict import TensorDict
>>> import torch
>>>
>>> tensordict_source = TensorDict({"a": torch.randn(3), "b": {"c": torch.randn(3)}}, [])
>>> state = {"state": tensordict}
>>> path = f"/tmp/{uuid.uuid4()}"
>>> snapshot = torchsnapshot.Snapshot.take(app_state=state, path=path)
>>> # later
>>> snapshot = torchsnapshot.Snapshot(path=path)
>>> tensordict2 = TensorDict()
>>> target_state = {
>>>     "state": tensordict2
>>> }
>>> snapshot.restore(app_state=target_state)
>>> assert (tensordict == tensordict2).all()

保存和加载大型数据集¶

如果数据集太大而无法放入内存，则上述方法可能会轻易失败。我们利用 torchsnapshot 的能力将张量分块加载到其预分配的目标位置。这要求您知道目标数据将具有的形状、设备等，但这是能够对模型或数据加载进行检查点检查的小代价：

与前面的示例相反，我们将不使用 `TensorDict` 的 `load_state_dict()` 方法，而是使用目标对象获得的 `state_dict`，并用保存的数据重新填充它。

同样，只需两行代码即可保存数据：

>>> app_state = {
...     "state": torchsnapshot.StateDict(tensordict=tensordict_source.state_dict(keep_vars=True))
... }
>>> snapshot = torchsnapshot.Snapshot.take(app_state=app_state, path="/path/to/my/snapshot")

我们使用了 `torchsnapshot.StateDict`，并且明确调用了 `my_tensordict_source.state_dict(keep_vars=True)`，这与前面的示例不同。现在，要将其加载到目标 tensordict 中：

>>> snapshot = Snapshot(path="/path/to/my/snapshot")
>>> app_state = {
...     "state": torchsnapshot.StateDict(tensordict=tensordict_target.state_dict(keep_vars=True))
... }
>>> snapshot.restore(app_state=app_state)

在此示例中，加载完全由 torchsnapshot 处理，即没有调用 `TensorDict.load_state_dict()`。

注意

这有两个重要含义：

由于 `LazyStackedTensorDict.state_dict()`（以及其他惰性 tensordict 类）在执行某些操作后返回数据的副本，因此加载到 state-dict 不会更新原始类。但是，由于支持 `state_dict()` 操作，这不会引发错误。
同样，由于 state-dict 是就地更新的，但 tensordict 没有使用 `TensorDict.update()` 或 `TensorDict.set()` 进行更新，因此目标 tensordict 中缺少键将不会被注意到。