pad_sequence

class tensordict.pad_sequence(list_of_tensordicts: Sequence[T], batch_first: bool | None = None, pad_dim: int = 0, padding_value: float = 0.0, out: T | None = None, device: DeviceType | None = None, return_mask: bool | NestedKey = False)

将一系列 tensordict 进行填充，以便它们能够以连续的格式堆叠在一起。

参数：

• list_of_tensordicts (List[TensorDictBase]) – 要填充和堆叠的实例列表。

• pad_dim (int, 可选) – pad_dim 指示要填充 tensordict 中所有键的维度。默认为 0。

• padding_value (数字, 可选) – 填充值。默认为 0.0。

• out (TensorDictBase, 可选) – 如果提供，则为数据将写入到的目标位置。

• return_mask (bool 或 NestedKey, 可选) – 如果为 True，则会返回“masks”条目。如果 return_mask 是嵌套键（字符串或字符串元组），它将返回掩码并用作掩码条目的键。它包含一个与堆叠的 tensordict 具有相同结构的 tensordict，其中每个条目都包含大小为 torch.Size([stack_len, *new_shape]) 的有效值的掩码，其中 new_shape[pad_dim] = max_seq_length，其余的 new_shape 与包含张量的先前形状匹配。

示例

>>> list_td = [
...     TensorDict({"a": torch.zeros((3, 8)), "b": torch.zeros((6, 8))}, batch_size=[]),
...     TensorDict({"a": torch.zeros((5, 8)), "b": torch.zeros((6, 8))}, batch_size=[]),
...     ]
>>> padded_td = pad_sequence(list_td, return_mask=True)
>>> print(padded_td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([2, 4, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([2, 5, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        masks: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.bool, is_shared=False),
                b: Tensor(shape=torch.Size([2, 6]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([2]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)