如何借助torch.unique移除形状为(n,m)的torch.tensor中的重复行？

torch.unique本身完全可以实现按行去重，之前理不清逻辑通常是踩了默认参数的坑：默认会把整个张量展平为一维，找元素级的唯一值，根本不会按整行做重复判断。

核心去重逻辑 #

针对形状为(n,m)的二维张量，行是第0维，只需要在调用torch.unique时指定dim=0，就会将每一行作为独立单元做重复比对，直接得到去重后的结果：

import torch

t1 = torch.tensor([[1, 2, 3], [4, 5, 6], [1, 2, 3], [4, 5, 6]])
# 沿行维度去重
t2 = torch.unique(t1, dim=0)
# 输出结果为 tensor([[1, 2, 3], [4, 5, 6]])

注意：默认参数下torch.unique会对去重结果按行字典序做排序，如果需要保留原张量中每一行第一次出现的相对顺序，需要传入sorted=False参数：

t2 = torch.unique(t1, dim=0, sorted=False)

定位重复行的位置 #

如果不需要直接得到去重后的张量，只想定位重复行的索引，可以利用torch.unique的返回参数实现：

• 传入return_inverse=True，会返回一个长度为n的数组，记录原始张量的每一行，对应哪个唯一行
• 传入return_counts=True，会返回每个唯一行在原始张量中的出现次数

示例代码：

unique_rows, inverse_idx, counts = torch.unique(
    t1, dim=0, return_inverse=True, return_counts=True
)
# 收集所有出现次数大于1的重复行的原始位置
duplicate_pos = torch.cat([
    torch.where(inverse_idx == row_id)[0]
    for row_id, cnt in enumerate(counts) if cnt > 1
])

如果需要实现“只保留首次出现的行、移除后续所有重复行”的逻辑，也可以基于逆映射数组快速生成保留掩码：

seen_ids = set()
keep_mask = []
for idx, row_id in enumerate(inverse_idx):
    if row_id.item() not in seen_ids:
        seen_ids.add(row_id.item())
        keep_mask.append(True)
    else:
        keep_mask.append(False)
# 按掩码索引提取保留的行
t2 = t1[torch.tensor(keep_mask)]

内容的提问来源于stack exchange，提问作者aretor