GRPOLossOutput¶

class torchrl.objectives.llm.GRPOLossOutput(loss_objective: 'torch.Tensor', clip_fraction: 'torch.Tensor', kl_approx: 'torch.Tensor', ESS: 'torch.Tensor', entropy: 'torch.Tensor | None' = None, loss_entropy: 'torch.Tensor | None' = None, loss_kl_to_ref: 'torch.Tensor | None' = None, kl_to_ref: 'torch.Tensor | None' = None, loss_kl_to_inference: 'torch.Tensor | None' = None, kl_to_inference: 'torch.Tensor | None' = None, *, batch_size, device=None, names=None)[源代码]¶

cat(dim: int = 0, *, out=None)¶

将tensordicts沿给定维度连接成一个tensordict。

此调用等同于调用 torch.cat()，但与 torch.compile 兼容。

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

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

将tensordict保存到磁盘。

此函数是 memmap() 的代理。

classmethod fields()¶

返回一个描述此数据类的字段的元组。字段类型为 Field。

接受一个数据类或其实例。元组元素为 Field 类型。

from_any(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)¶

Recursively converts any object to a TensorDict.

注意

from_any 比常规的 TensorDict 构造函数限制更少。它可以使用自定义启发式方法将数据结构（如 dataclasses 或 tuples）转换为 tensordict。此方法可能会产生一些额外的开销，并在映射策略方面涉及更多主观选择。

注意

This method recursively converts the input object to a TensorDict. If the object is already a TensorDict (or any similar tensor collection object), it will be returned as is.

参数:

obj – The object to be converted.

关键字参数:

auto_batch_size (bool, optional) – 如果 True，将自动计算 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应具有的维度数。默认为 None（每个级别的完整 batch size）。
device (torch.device, optional) – 将创建 TensorDict 的设备。
batch_size (torch.Size, optional) – TensorDict 的批次大小。与 auto_batch_size 互斥。

返回:

A TensorDict representation of the input object.

Supported objects

通过 from_dataclass() 创建的数据类（数据类将转换为 TensorDict 实例，而非 tensorclass）。
通过 from_namedtuple() 创建的命名元组。
通过 from_dict() 的字典。
通过 from_tuple() 创建的元组。
通过 from_struct_array() 创建的 NumPy 结构化数组。
通过 from_h5() 创建的 HDF5 对象。

from_dataclass(*, dest_cls: Type | None = None, auto_batch_size: bool = False, batch_dims: int | None = None, as_tensorclass: bool = False, device: torch.device | None = None, batch_size: torch.Size | None = None)¶

Converts a dataclass into a TensorDict instance.

参数:

dataclass – The dataclass instance to be converted.

关键字参数:

dest_cls (tensorclass, optional) – 用于映射数据的 tensorclass 类型。如果未提供，则创建一个新类。如果 obj 是一个类型或 as_tensorclass 为 False，则无效。
auto_batch_size (bool, optional) – 如果 True，将自动确定并应用 batch size 到生成的 TensorDict。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应具有的维度数。默认为 None（每个级别的完整 batch size）。
as_tensorclass (bool, optional) – 如果 True，则将转换委托给自由函数 from_dataclass()，并返回一个 tensorclass（tensorclass()）类型或实例，而不是 TensorDict。默认为 False。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

A TensorDict instance derived from the provided dataclass, unless as_tensorclass is True, in which case a tensor-compatible class or instance is returned.

抛出:

TypeError – 如果提供的输入不是数据类实例。

警告

This method is distinct from the free function from_dataclass and serves a different purpose. While the free function returns a tensor-compatible class or instance, this method returns a TensorDict instance.

注意

此方法创建一个新的 TensorDict 实例，其键对应于输入 dataclass 的字段。
结果 TensorDict 中的每个键都使用 `cls.from_any` 方法进行初始化。
auto_batch_size 选项允许自动确定批次大小并将其应用于结果 TensorDict。

from_h5(*, mode: str = 'r', auto_batch_size: bool = False, batch_dims: int | None = None, batch_size: torch.Size | None = None)¶

从 h5 文件创建 PersistentTensorDict。

参数:

filename (str) – h5 文件的路径。

关键字参数:

mode (str, optional) – 读取模式。默认为 "r"。
auto_batch_size (bool, optional) – 如果 True，将自动计算 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应具有的维度数。默认为 None（每个级别的完整 batch size）。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入 h5 文件的 PersistentTensorDict 表示。

示例

>>> td = TensorDict.from_h5("path/to/file.h5")
>>> print(td)
PersistentTensorDict(
    fields={
        key1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        key2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

from_modules(*, as_module: bool = False, lock: bool = True, use_state_dict: bool = False, lazy_stack: bool = False, expand_identical: bool = False)¶

为 vmap 的 ensemable 学习/特征期望应用检索多个模块的参数。

参数:

modules (nn.Module 序列) – 要从中获取参数的模块。如果模块结构不同，则需要懒惰堆叠（请参阅下面的 `lazy_stack` 参数）。

关键字参数:

as_module (bool, optional) – 如果 True，则返回一个 TensorDictParams 实例，可用于将参数存储在 torch.nn.Module 中。默认为 False。
lock (bool, optional) – 如果 True，则锁定的结果 tensordict。默认为 True。
use_state_dict (bool, optional) –
如果为 True，则将使用模块的 state-dict 并将其解压到具有模型树结构的 TensorDict 中。默认为 False。

注意

这在使用 state-dict hook 时尤其有用。
lazy_stack (bool, optional) –
是否密集堆叠或懒惰堆叠参数。默认为 False（密集堆叠）。

注意

lazy_stack 和 as_module 是互斥的特性。

警告

懒惰输出和非懒惰输出之间有一个重要的区别：非懒惰输出将使用所需的批次大小重新实例化参数，而 `lazy_stack` 将仅将参数表示为懒惰堆叠。这意味着，虽然原始参数可以安全地传递给优化器（当 `lazy_stack=True` 时），但在设置为 `True` 时需要传递新参数。

警告

虽然使用 lazy stack 来保持原始参数引用可能很诱人，但请记住，每次调用 get() 时，lazy stack 都会执行一个堆栈。这需要内存（参数大小的 N 倍，如果构建了图，则更多）和计算时间。这也意味着优化器将包含更多参数，像 step() 或 zero_grad() 这样的操作的执行时间会更长。总的来说，lazy_stack 应保留给极少数用例。
expand_identical (bool, optional) – 如果 True 且相同的参数（相同标识）被堆叠到自身，则将返回该参数的扩展版本。在 lazy_stack=True 时忽略此参数。

示例

>>> from torch import nn
>>> from tensordict import TensorDict
>>> torch.manual_seed(0)
>>> empty_module = nn.Linear(3, 4, device="meta")
>>> n_models = 2
>>> modules = [nn.Linear(3, 4) for _ in range(n_models)]
>>> params = TensorDict.from_modules(*modules)
>>> print(params)
TensorDict(
    fields={
        bias: Parameter(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Parameter(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)
>>> # example of batch execution
>>> def exec_module(params, x):
...     with params.to_module(empty_module):
...         return empty_module(x)
>>> x = torch.randn(3)
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> assert y.shape == (n_models, 4)
>>> # since lazy_stack = False, backprop leaves the original params untouched
>>> y.sum().backward()
>>> assert params["weight"].grad.norm() > 0
>>> assert modules[0].weight.grad is None

当 lazy_stack=True 时，情况略有不同

>>> params = TensorDict.from_modules(*modules, lazy_stack=True)
>>> print(params)
LazyStackedTensorDict(
    fields={
        bias: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Tensor(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    exclusive_fields={
    },
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False,
    stack_dim=0)
>>> # example of batch execution
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> assert y.shape == (n_models, 4)
>>> y.sum().backward()
>>> assert modules[0].weight.grad is not None

from_namedtuple(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)¶

递归地将命名元组转换为 TensorDict。

参数:

named_tuple – 要转换的命名元组实例。

关键字参数:

auto_batch_size (bool, optional) – 如果 True，将自动计算 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应具有的维度数。默认为 None（每个级别的完整 batch size）。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入命名元组的 TensorDict 表示。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> nt = data.to_namedtuple()
>>> print(nt)
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))
>>> TensorDict.from_namedtuple(nt, auto_batch_size=True)
TensorDict(
    fields={
        a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                a_string: NonTensorData(data=zero!, batch_size=torch.Size([3]), device=None),
                a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

from_pytree(*, batch_size: torch.Size | None = None, auto_batch_size: bool = False, batch_dims: int | None = None)¶

将 pytree 转换为 TensorDict 实例。

此方法旨在尽可能保留 pytree 的嵌套结构。

其他非张量键将被添加，以跟踪每个级别的标识，从而提供内置的 pytree 到 tensordict 的双射转换 API。

当前接受的类包括列表、元组、命名元组和字典。

注意

对于字典，非 NestedKey 键会作为 NonTensorData 实例单独注册。

注意

可转换为张量类型（如 int、float 或 np.ndarray）将被转换为 torch.Tensor 实例。请注意，此转换是满射的：将 tensordict 转换回 pytree 将无法恢复原始类型。

示例

>>> # Create a pytree with tensor leaves, and one "weird"-looking dict key
>>> class WeirdLookingClass:
...     pass
...
>>> weird_key = WeirdLookingClass()
>>> # Make a pytree with tuple, lists, dict and namedtuple
>>> pytree = (
...     [torch.randint(10, (3,)), torch.zeros(2)],
...     {
...         "tensor": torch.randn(
...             2,
...         ),
...         "td": TensorDict({"one": 1}),
...         weird_key: torch.randint(10, (2,)),
...         "list": [1, 2, 3],
...     },
...     {"named_tuple": TensorDict({"two": torch.ones(1) * 2}).to_namedtuple()},
... )
>>> # Build a TensorDict from that pytree
>>> td = TensorDict.from_pytree(pytree)
>>> # Recover the pytree
>>> pytree_recon = td.to_pytree()
>>> # Check that the leaves match
>>> def check(v1, v2):
>>>     assert (v1 == v2).all()
>>>
>>> torch.utils._pytree.tree_map(check, pytree, pytree_recon)
>>> assert weird_key in pytree_recon[1]

from_remote_init(group: 'ProcessGroup' | None = None, device: torch.device | None = None) → Self¶

从远程发送的元数据创建新的 tensordict 实例。

此类方法接收由 init_remote 发送的元数据，创建具有匹配形状和 dtype 的新 tensordict，然后异步接收实际的 tensordict 内容。

参数:

src (int) – 发送元数据的源进程的 rank。
group ("ProcessGroup", optional) – 要使用的进程组。默认为 None。
device (torch.device, 可选) – 用于张量运算的设备。默认为 None。

返回:

使用接收到的元数据和内容初始化的新 tensordict 实例。

返回类型:

TensorDict

另请参阅

发送进程应已调用 ~.init_remote 来发送元数据和内容。

from_struct_array(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None) → Self¶

将结构化 numpy 数组转换为 TensorDict。

生成的 TensorDict 将与 numpy 数组共享相同的内存内容（这是一次零拷贝操作）。原地更改结构化 numpy 数组的值会影响 TensorDict 的内容。

注意

此方法执行零拷贝操作，这意味着生成的 TensorDict 将与输入的 numpy 数组共享相同的内存内容。因此，原地更改 numpy 数组的值会影响 TensorDict 的内容。

参数:

struct_array (np.ndarray) – 要转换的结构化 numpy 数组。

关键字参数:

auto_batch_size (bool, optional) – 如果 True，将自动计算 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应具有的维度数。默认为 None（每个级别的完整 batch size）。
device (torch.device, 可选) –
将创建 TensorDict 的设备。默认为 None。

注意

更改设备（即，指定任何非 `None` 或 `"cpu"` 的设备）将传输数据，从而导致返回数据的内存位置发生更改。
batch_size (torch.Size, 可选) – TensorDict 的批次大小。默认为 None。

返回:

输入的结构化 numpy 数组的 TensorDict 表示。

示例

>>> x = np.array(
...     [("Rex", 9, 81.0), ("Fido", 3, 27.0)],
...     dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")],
... )
>>> td = TensorDict.from_struct_array(x)
>>> x_recon = td.to_struct_array()
>>> assert (x_recon == x).all()
>>> assert x_recon.shape == x.shape
>>> # Try modifying x age field and check effect on td
>>> x["age"] += 1
>>> assert (td["age"] == np.array([10, 4])).all()

classmethod from_tensordict(tensordict: TensorDictBase, non_tensordict: dict | None = None, safe: bool = True) → Self¶

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

参数:

tensordict (TensorDictBase) – 张量类型字典
non_tensordict (dict) – 包含非张量和嵌套张量类对象的字典
safe (bool) – 如果 tensordict 不是 TensorDictBase 实例，是否引发错误

from_tuple(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)¶

将元组转换为 TensorDict。

参数:

obj – 要转换的元组实例。

关键字参数:

auto_batch_size (bool, optional) – 如果 True，将自动计算 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应具有的维度数。默认为 None（每个级别的完整 batch size）。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入的元组的 TensorDict 表示。

示例

>>> my_tuple = (1, 2, 3)
>>> td = TensorDict.from_tuple(my_tuple)
>>> print(td)
TensorDict(
    fields={
        0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

fromkeys(value: Any = 0)¶

从键列表和单个值创建 tensordict。

参数:

keys (list of NestedKey) – 指定新字典键的可迭代对象。
value (compatible type, optional) – 所有键的值。默认为 0。

get(key: NestedKey, *args, **kwargs)¶

获取输入键对应的存储值。

参数:

key (str, tuple of str) – 要查询的键。如果是字符串元组，则等同于链式调用 getattr。
default – 如果在张量类中找不到键，则返回默认值。

返回:

存储在输入键下的值

lazy_stack(dim: int = 0, *, out=None, **kwargs)¶

创建 TensorDicts 的懒惰堆叠。

有关详细信息，请参阅 lazy_stack()。

load(*args, **kwargs) → Self¶

从磁盘加载 tensordict。

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

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

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

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

load_memmap(device: torch.device | None = None, non_blocking: bool = False, *, out: TensorDictBase | None = None) → Self¶

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

参数:

prefix (str 或 Path to folder) – 应获取已保存 tensordict 的文件夹路径。
device (torch.device 或 等效项, 可选) – 如果提供，数据将异步转换为该设备。支持 `"meta"` 设备，在这种情况下，数据不会被加载，而是创建一组空的 "meta" 张量。这对于在不实际打开任何文件的情况下了解模型大小和结构很有用。
non_blocking (bool, optional) – 如果为 True，则在将张量加载到设备后不会调用同步。默认为 False。
out (TensorDictBase, optional) – 数据应写入的可选 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()

此方法还允许加载嵌套的 tensordicts。

示例

>>> 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_dense_stack(dim: int = 0, *, out=None, **kwargs)¶

尝试使 TensorDicts 密集堆叠，并在需要时回退到懒惰堆叠。

有关详细信息，请参阅 maybe_dense_stack()。

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

将所有张量写入内存映射的 Tensor 中，并放入新的 tensordict。

参数:

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 的未来对象。
share_non_tensor (bool, optional) – 如果为 True，则非张量数据将在进程之间共享，并且单个节点内任何工作进程上的写入操作（例如原地更新或设置）将更新所有其他工作进程上的值。如果非张量叶子节点数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, optional) – 如果为 False，则如果同一路径下已存在张量，将引发异常。默认为 True。

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

返回:: 返回一个新的 tensordict，其中张量存储在磁盘上（如果 return_early=False），否则返回一个 TensorDictFuture 实例。

注意

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

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

将所有张量原地写入相应的内存映射张量。

参数:

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.result() 查询结果 tensordict。
share_non_tensor (bool, optional) – 如果为 True，则非张量数据将在进程之间共享，并且单个节点内任何工作进程上的写入操作（例如原地更新或设置）将更新所有其他工作进程上的值。如果非张量叶子节点数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, optional) – 如果为 False，则如果同一路径下已存在张量，将引发异常。默认为 True。

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

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

注意

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

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

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

参数:

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 的未来对象。
share_non_tensor (bool, optional) – 如果为 True，则非张量数据将在进程之间共享，并且单个节点内任何工作进程上的写入操作（例如原地更新或设置）将更新所有其他工作进程上的值。如果非张量叶子节点数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, optional) – 如果为 False，则如果同一路径下已存在张量，将引发异常。默认为 True。

然后，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，则刷新其内容。

如果没有任何路径与之关联，此方法将引发异常。

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

将tensordict保存到磁盘。

此函数是 memmap() 的代理。

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

设置一个新的键值对。

参数:

key (str, tuple of str) – 要设置的键的名称。如果是字符串元组，则等同于链式调用 getattr，然后进行最终的 setattr。
value (Any) – 要存储在张量类中的值
inplace (bool, optional) – 如果为 True，则 set 将尝试原地更新值。如果为 False 或键不存在，则值将简单地写入其目的地。

返回:

self

stack(dim: int = 0, *, out=None)¶

沿给定维度将 tensordicts 堆叠成一个单一的 tensordict。

此调用等效于调用 torch.stack()，但与 torch.compile 兼容。

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

to_tensordict(*, retain_none: bool | None = None) → TensorDict¶

将张量类转换为常规 TensorDict。

复制所有条目。内存映射和共享内存张量将被转换为常规张量。

参数:: retain_none (bool) – 如果为 True，则 None 值将被写入 tensordict。否则它们将被丢弃。默认值：True。
返回:: 包含与张量类相同值的新的 TensorDict 对象。

unbind(dim: int)¶

返回沿指定维度解绑的索引张量类实例的元组。

结果张量类实例将共享初始张量类实例的存储。

GRPOLossOutput¶

文档

教程

资源