remove_duplicates

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

移除 key 中沿指定維度重複的索引。

此方法檢測與指定 key 相關聯的 tensor 中沿指定 dim 的重複元素，並移除 TensorDict 中所有其他 tensor 中相同索引中的元素。 預期 dim 是輸入 TensorDict 批次大小中的其中一個維度，以確保所有 tensor 中的一致性。 否則，將引發錯誤。

參數:

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

• key (NestedKey) – 應該識別並移除重複元素的 tensor 的鍵。 它必須是 TensorDict 中的其中一個葉鍵，指向一個 tensor，而不是另一個 TensorDict。

• dim (int, optional) – 應該識別並移除重複元素的維度。 它必須是輸入 TensorDict 批次大小中的其中一個維度。 預設為 0。

• return_indices (bool, optional) – 如果為 True，也會傳回輸入 tensor 中唯一元素的索引。 預設為 False。

回傳:

輸入 tensordict，其中包含在 dimension dim 中

key tensor 內重複元素對應的索引已被移除。

unique_indices (torch.Tensor, optional): 輸入 tensordict 中

指定 key 沿指定 dim 的第一個唯一元素出現位置的索引。 僅在 return_index 為 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()