remove_duplicates

class tensordict.utils.remove_duplicates(input: TensorDictBase, key: NestedKey, dim: int = 0, *, return_indices: bool = False)

移除指定维度 key 中的重复索引。

此方法检测与指定的 key 关联的张量沿指定的 dim 的重复元素，并移除 TensorDict 内所有其他张量中相同索引处的元素。预期 dim 是输入 TensorDict 批大小内的维度之一，以确保所有张量的一致性。否则，将引发错误。

参数:

• input (TensorDictBase) – 包含可能重复元素的 TensorDict。

• key (NestedKey) – 用于标识和移除重复元素的张量的键。它必须是 TensorDict 中的叶键之一，指向张量而不是另一个 TensorDict。

• dim (int, optional) – 用于标识和移除重复元素的维度。它必须是输入 TensorDict 批大小内的维度之一。默认为 0。

• return_indices (bool, optional) – 如果为 True，则还将返回输入张量中唯一元素的索引。默认为 False。

返回值:

移除了张量 key 沿维度 dim 的重复元素的索引的输入 tensordict。

在张量 key 沿维度 dim 中，唯一元素的首次出现的索引 (torch.Tensor，可选)。仅当 return_index 为 True 时提供。

unique_indices (torch.Tensor, optional):

input tensordict for the specified key along the specified dim. Only provided if return_index is True.

返回类型:

output (TensorDictBase)

示例

>>> td = TensorDict(
...     {
...         "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [7, 8, 9]]),
...         "tensor2": torch.tensor([[10, 20], [30, 40], [40, 50], [50, 60]]),
...     }
...     batch_size=[4],
... )
>>> output_tensordict = remove_duplicate_elements(td, key="tensor1", dim=0)
>>> expected_output = TensorDict(
...     {
...         "tensor1": torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
...         "tensor2": torch.tensor([[10, 20], [30, 40], [50, 60]]),
...     },
...     batch_size=[3],
... )
>>> assert (td == expected_output).all()