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)¶
将所有张量写入内存映射的 Tensor 中,并放入新的 tensordict。
- 参数:
data (TensorDictBase) – 要保存的数据结构。
prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。
copy_existing (bool) – 如果为 False(默认值),如果 tensordict 中的某个条目已经是存储在磁盘上的张量且具有关联文件,但未根据 prefix 保存到正确的位置,则会引发异常。如果为
True
,则会将任何现有张量复制到新位置。
- 关键字参数:
num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0。
return_early (bool, optional) – 如果为
True
且num_threads>0
,则该方法将返回一个 tensordict 的 future。share_non_tensor (bool, optional) – 如果为
True
,则非张量数据将在进程之间共享,并且在单个节点内的任何工作进程上进行的写入操作(如原地更新或设置)将更新所有其他工作进程上的值。如果非张量叶子节点数量很高(例如,共享大量非张量数据),这可能会导致 OOM 或类似错误。默认为False
。
然后,TensorDict 会被锁定,这意味着任何非原地写入操作都会引发异常(例如,重命名、设置或移除条目)。一旦 tensordict 被解锁,内存映射属性将变为
False
,因为不能保证跨进程的同一性。- 返回:
如果
return_early=False
,则返回一个包含已存储在磁盘上的张量的新 tensordict;否则,返回一个TensorDictFuture
实例。
注意
以这种方式序列化对于深度嵌套的 tensordicts 来说可能很慢,因此不建议在训练循环中调用此方法。