【PyTorch】torch.sort() 函数：对张量进行排序

torch.sort 函数

torch.sort 是 PyTorch 中用于对张量进行排序的函数。它会返回一个排序后的张量，并根据指定的维度对张量进行排序，支持升序和降序排序。

函数签名：

torch.sort(input, dim=-1, descending=False, out=None)

参数说明：

• input (Tensor): 需要排序的输入张量。
• dim (int, optional): 排序操作沿着哪个维度进行。默认值是 -1，表示按最后一个维度进行排序。如果是 0，则按第一个维度排序。
• descending (bool, optional): 如果设置为 True，则对张量进行降序排序；默认值是 False，表示升序排序。
• out (tuple, optional): 一个元组 (values, indices)，用于存储返回结果。values 是排序后的张量，indices 是排序过程中元素在原张量中的索引位置。

返回值：

• values: 排序后的张量。
• indices: 对应排序后元素在原张量中的索引位置。

示例：

1. 对一维张量进行排序：

import torch

# 创建一个一维张量
tensor = torch.tensor([3, 1, 4, 1, 5, 9, 2])

# 升序排序
sorted_tensor, indices = torch.sort(tensor)
print("Sorted Tensor:", sorted_tensor)
print("Indices:", indices)

输出：

Sorted Tensor: tensor([1, 1, 2, 3, 4, 5, 9])
Indices: tensor([1, 3, 6, 0, 2, 4, 5])

在这个例子中，sorted_tensor 是排序后的结果，indices 是每个元素在原张量中的索引位置。

2. 对二维张量按列进行排序：

import torch

# 创建一个二维张量
tensor = torch.tensor([[3, 1, 4], [1, 5, 9], [2, 6, 5]])

# 按列（dim=0）升序排序
sorted_tensor, indices = torch.sort(tensor, dim=0)
print("Sorted Tensor (by columns):\n", sorted_tensor)
print("Indices:\n", indices)

输出：

Sorted Tensor (by columns):
 tensor([[1, 1, 4],
         [2, 5, 5],
         [3, 6, 9]])
Indices:
 tensor([[1, 0, 0],
         [2, 2, 1],
         [0, 1, 2]])

此时，sorted_tensor 是按列升序排序后的张量，indices 是每个元素在原张量中的行索引。

3. 对二维张量按行进行排序：

import torch

# 创建一个二维张量
tensor = torch.tensor([[3, 1, 4], [1, 5, 9], [2, 6, 5]])

# 按行（dim=1）升序排序
sorted_tensor, indices = torch.sort(tensor, dim=1)
print("Sorted Tensor (by rows):\n", sorted_tensor)
print("Indices:\n", indices)

输出：

Sorted Tensor (by rows):
 tensor([[1, 3, 4],
         [1, 5, 9],
         [2, 5, 6]])
Indices:
 tensor([[1, 0, 2],
         [0, 1, 2],
         [0, 2, 1]])

这里，sorted_tensor 是按行升序排序后的张量，indices 是每个元素在原张量中的列索引。

4. 使用降序排序：

import torch

# 创建一个一维张量
tensor = torch.tensor([3, 1, 4, 1, 5, 9, 2])

# 降序排序
sorted_tensor, indices = torch.sort(tensor, descending=True)
print("Sorted Tensor (Descending):", sorted_tensor)
print("Indices:", indices)

输出：

Sorted Tensor (Descending): tensor([9, 5, 4, 3, 2, 1, 1])
Indices: tensor([5, 4, 2, 0, 6, 1, 3])

5. 排序并返回 indices：

有时你只需要排序后的索引，而不关心排序的值。你可以通过只返回 indices 来获得元素的排序顺序。

import torch

# 创建一个一维张量
tensor = torch.tensor([3, 1, 4, 1, 5, 9, 2])

# 只返回索引
_, indices = torch.sort(tensor)
print("Indices:", indices)

输出：

Indices: tensor([1, 3, 6, 0, 2, 4, 5])

常见用途：

1. 排序数据：在处理排序问题时，torch.sort 很有用，比如对数据进行排序后再进行后续处理。
2. 获取最大/最小值的索引：通过排序和索引，你可以轻松地获得数据中最大或最小值的位置。
3. 排序神经网络的输出：在一些任务中，如排序学习或推荐系统，可能需要对神经网络的输出进行排序。
4. 进行排名：通过获取排序的索引，您可以生成排名信息，常用于信息检索和推荐系统等领域。

---

总结：

torch.sort 是一个非常有用的函数，可以对张量进行排序。它返回排序后的张量以及排序时对应的索引。在多维张量中，可以选择排序的维度和排序方式（升序或降序）。这个函数对于数据处理、特征排序和计算排名等任务非常有帮助。