快捷方式

memmap

class tensordict.memmap(data: TensorDictBase, prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False)

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

参数:
  • data (TensorDictBase) – 要保存的数据结构。

  • prefix (str) – 内存映射张量将要存储的目录前缀。目录树结构将模仿 tensordict 的结构。

  • copy_existing (bool) – 如果为 False(默认),如果 tensordict 中的条目已经是存储在磁盘上的张量并关联了文件,但未根据前缀保存在正确的位置,则会引发异常。如果为 True,则任何现有张量都将被复制到新位置。

关键词参数:
  • num_threads (int, 可选) – 用于写入 memmap 张量的线程数。默认为 0

  • return_early (bool, 可选) – 如果 Truenum_threads>0,则该方法将返回 tensordict 的 future。

  • share_non_tensor (bool, 可选) – 如果 True,非张量数据将在进程和写入操作之间共享(例如,在单个节点内的任何 worker 上进行原地更新或设置都将更新所有其他 worker 上的值)。如果非张量叶节点的数量很高(例如,共享大量非张量数据堆栈),则可能导致 OOM 或类似错误。默认为 False

然后 TensorDict 会被锁定,这意味着任何非原地写入操作都将抛出异常(例如,重命名、设置或删除条目)。一旦 tensordict 被解锁,内存映射属性将变为 False,因为不再保证跨进程的同一性。

返回:

如果 return_early=False,则返回一个新的张量存储在磁盘上的 tensordict,否则返回 TensorDictFuture 实例。

注意

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得您的问题解答

查看资源