如何获取二维tensor中取值相同的行张量对应的索引分组？

给定你的输入张量：

import torch

input_tensor = torch.tensor([
    [0, 1, 0, 1], # A
    [1, 1, 0, 1], # B
    [1, 0, 0, 1], # C
    [0, 1, 0, 1], # D
    [1, 1, 0, 1], # E
    [1, 1, 0, 1]  # F
])

实现代码 #

直接用PyTorch自带的unique接口就能高效完成，不需要手动写循环匹配：

# 按行去重，同时返回每个原始行对应去重后行的索引标签
_, label = torch.unique(input_tensor, dim=0, return_inverse=True)

# 按标签分组，过滤掉只有单个元素的分组（即无重复的C行）
result = []
for class_id in label.unique():
    idx_list = torch.where(label == class_id)[0].tolist()
    if len(idx_list) > 1:
        result.append(tuple(idx_list))

执行后result的输出就是你要求的[(0, 3), (1, 4, 5)]。

补充说明 #

如果你需要保留所有行的分组（包括只有C行的(2,)分组），删掉if len(idx_list) > 1的判断即可。

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