如何基于指定列对PyTorch二维Tensor的行进行排序？

原始二维Tensor：

import torch
original_tensor = torch.tensor([[0.8771, 0.0976, 0.8186],
                                [0.7044, 0.4783, 0.0350],
                                [0.4239, 0.8341, 0.3693],
                                [0.5568, 0.9175, 0.0763],
                                [0.0876, 0.1651, 0.2776]])

需求：按指定列（比如最后一列）的数值升序排列所有行，得到如下结果：

tensor([[0.7044, 0.4783, 0.0350],
        [0.5568, 0.9175, 0.0763],
        [0.0876, 0.1651, 0.2776],
        [0.4239, 0.8341, 0.3693],
        [0.8771, 0.0976, 0.8186]])

具体实现步骤 #

• 提取用于排序的目标列：比如最后一列用original_tensor[:, -1]，第二列用original_tensor[:, 1]（索引从0开始计数）
• 调用torch.sort()获取该列排序后的行索引，这个索引就是原Tensor行的重排顺序
• 用索引对原Tensor进行切片，得到排序后的结果

完整代码示例 #

import torch

original_tensor = torch.tensor([[0.8771, 0.0976, 0.8186],
                                [0.7044, 0.4783, 0.0350],
                                [0.4239, 0.8341, 0.3693],
                                [0.5568, 0.9175, 0.0763],
                                [0.0876, 0.1651, 0.2776]])

# 指定要排序的列，这里选最后一列（索引-1）
target_col = -1
# 获取排序后的行索引，sorted_values是排序后的列值，这里我们主要用sorted_indices
_, sorted_indices = torch.sort(original_tensor[:, target_col])

# 按索引重排行，得到最终排序后的Tensor
sorted_tensor = original_tensor[sorted_indices]

print(sorted_tensor)

额外说明 #

如果需要降序排序，只需要在torch.sort()里添加descending=True参数：

_, sorted_indices_desc = torch.sort(original_tensor[:, target_col], descending=True)
sorted_tensor_desc = original_tensor[sorted_indices_desc]

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