Torchsort

在 PyTorch 中实现快速、可微分的排序和排名。

这是 Fast Differentiable Sorting and Ranking（Blondel 等人）的纯 PyTorch 实现。大部分代码直接复制自原始的 NumPy 实现，地址为 google-research/fast-soft-sort，其中等序回归求解器被重写为 PyTorch C++ 和 CUDA 扩展。

安装

pip install torchsort

要构建 CUDA 扩展，您需要安装 CUDA 工具链。如果您希望在没有 CUDA 运行时的环境中（例如 Docker）进行构建，则需要在安装之前导出环境变量 TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing;Ampere"。

Conda 安装

在某些系统上，使用 `pip` 在 Conda 环境中安装该包可能会编译失败。如果发生这种情况，您可能需要：

1. 使用 conda install -c conda-forge gxx_linux-64=9.40 安装 g++。
2. 运行 export CXX=/path/to/miniconda3/envs/env_name/bin/x86_64-conda_cos6-linux-gnu-g++。
3. 运行 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/miniconda3/lib。
4. 运行 pip install --force-reinstall --no-cache-dir --no-deps torchsort。

感谢 @levnikmyskin 和 @sachit-menon 指出这一问题！

预编译的轮子包

目前，Linux 上针对近期的 Python/PyTorch/CUDA 组合提供了预编译的轮子包：

# torchsort 版本，支持 >= 0.1.10
export TORCHSORT=0.1.10
# PyTorch 版本，分别支持 pt26、pt25、pt24、pt21、pt20 和 pt113，对应版本 2.6、2.5、2.4、2.1、2.0 和 1.13
export TORCH=pt26
# CUDA 版本，支持 cpu、cu113、cu117、cu118、cu121、cu124 和 cu126，分别对应仅 CPU、CUDA 11.3、CUDA 11.7、CUDA 11.8、CUDA 12.1、CUDA 12.4 和 CUDA 12.6
export CUDA=cu126
# Python 版本，支持 cp310、cp311 和 cp312，分别对应版本 3.10、3.11 和 3.12
export PYTHON=cp312

pip install https://github.com/teddykoker/torchsort/releases/download/v${TORCHSORT}/torchsort-${TORCHSORT}+${TORCH}${CUDA}-${PYTHON}-${PYTHON}-linux_x86_64.whl

感谢 siddharthab 帮助创建了构建流程！请参阅 最新发布 的 Assets 部分，以获取受支持组合的列表。

使用方法

torchsort 提供两个函数：soft_rank 和 soft_sort，每个函数都有 regularization（"l2" 或 "kl"）和 regularization_strength（标量值）两个参数。它们都会对二维张量的最后一维进行排名或排序，其准确性取决于正则化强度：

import torch
import torchsort

x = torch.tensor([[8, 0, 5, 3, 2, 1, 6, 7, 9]])

torchsort.soft_sort(x, regularization_strength=1.0)
# tensor([[0.5556, 1.5556, 2.5556, 3.5556, 4.5556, 5.5556, 6.5556, 7.5556, 8.5556]])
torchsort.soft_sort(x, regularization_strength=0.1)
# tensor([[-0., 1., 2., 3., 5., 6., 7., 8., 9.]])

torchsort.soft_rank(x)
# tensor([[8., 1., 5., 4., 3., 2., 6., 7., 9.]])

这两种操作在 CPU 或 GPU 上都是完全可微的：

x = torch.tensor([[8., 0., 5., 3., 2., 1., 6., 7., 9.]], requires_grad=True).cuda()
y = torchsort.soft_sort(x)

torch.autograd.grad(y[0, 0], x)
# (tensor([[0.1111, 0.1111, 0.1111, 0.1111, 0.1111, 0.1111, 0.1111, 0.1111, 0.1111]],
#         device='cuda:0'),)

示例

斯皮尔曼等级相关系数

Spearman's rank coefficient 是一种非常有用的指标，用于衡量两个变量之间的单调相关性。我们可以使用 Torchsort 创建一个可微分的斯皮尔曼等级相关系数函数，以便直接针对该指标优化模型：

import torch
import torchsort

def spearmanr(pred, target, **kw):
    pred = torchsort.soft_rank(pred, **kw)
    target = torchsort.soft_rank(target, **kw)
    pred = pred - pred.mean()
    pred = pred / pred.norm()
    target = target - target.mean()
    target = target / target.norm()
    return (pred * target).sum()

pred = torch.tensor([[1., 2., 3., 4., 5.]], requires_grad=True)
target = torch.tensor([[5., 6., 7., 8., 7.]])
spearman = spearmanr(pred, target)
# tensor(0.8321)

torch.autograd.grad(spearman, pred)
# (tensor([[-5.5470e-02,  2.9802e-09,  5.5470e-02,  1.1094e-01, -1.1094e-01]]),)

基准测试

torchsort 和 fast_soft_sort 的时间复杂度均为 O(n log n)，但与内置的 torch.sort 相比，两者都有一些额外的开销。当批次大小为 1 时（见左侧），fast_soft_sort 的 Numba JIT 编译前向传递性能与 torchsort 的 CPU 内核大致相当，然而其反向传播仍然依赖于部分 Python 代码，这大大降低了其性能。

此外，torchsort 的内核支持批量处理，并且随着批次大小的增加，其性能明显优于 fast_soft_sort。

torchsort 的 CUDA 内核在序列长度低于 ~2000 时表现良好，并且可以扩展到非常大的批次规模。未来，CUDA 内核还有望进一步优化，以达到接近内置 torch.sort 的性能。

参考文献

@inproceedings{blondel2020fast,
  title={Fast differentiable sorting and ranking},
  author={Blondel, Mathieu and Teboul, Olivier and Berthet, Quentin and Djolonga, Josip},
  booktitle={International Conference on Machine Learning},
  pages={950--959},
  year={2020},
  organization={PMLR}
}

Torchsort 快速上手指南

Torchsort 是一个基于 PyTorch 的高性能、可微分排序与排名工具库。它实现了快速可微排序算法，支持 CPU 和 GPU 加速，适用于需要优化排序指标（如 Spearman 相关系数）的深度学习场景。

环境准备

• 操作系统: Linux (预编译包主要支持 Linux)，Windows/macOS 需本地编译。
• Python: 推荐 3.10, 3.11 或 3.12。
• PyTorch: 支持主流版本 (1.13 - 2.6)。
• CUDA (可选): 若需使用 GPU 加速，需安装对应的 CUDA Toolkit。
  • 若在无 CUDA 运行时的环境（如 Docker）中构建，需设置环境变量：

    export TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing;Ampere"
    

• 编译器: 本地编译需安装 g++。在 Conda 环境中若编译失败，可能需要安装 gxx_linux-64 并配置 CXX 和 LD_LIBRARY_PATH。

安装步骤

方式一：直接安装 (推荐)

大多数情况下，直接使用 pip 安装即可自动获取预编译包或触发本地编译：

pip install torchsort

国内加速建议：如遇下载缓慢，可使用清华或阿里镜像源：

pip install torchsort -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：手动安装预编译 Wheel (Linux)

如果自动安装失败，可根据你的环境组合手动下载指定版本的 Wheel 包。以下示例展示了如何构建下载链接（需替换版本号）：

# 设置版本变量 (示例：torchsort 0.1.10, PyTorch 2.6, CUDA 12.6, Python 3.12)
export TORCHSORT=0.1.10
export TORCH=pt26
export CUDA=cu126
export PYTHON=cp312

# 安装指定 wheel
pip install https://github.com/teddykoker/torchsort/releases/download/v${TORCHSORT}/torchsort-${TORCHSORT}+${TORCH}${CUDA}-${PYTHON}-${PYTHON}-linux_x86_64.whl

基本使用

Torchsort 提供两个核心函数：soft_sort (软排序) 和 soft_rank (软排名)。它们对张量的最后一个维度进行操作，且完全支持自动求导。

1. 基础排序与排名

import torch
import torchsort

# 创建输入张量
x = torch.tensor([[8, 0, 5, 3, 2, 1, 6, 7, 9]])

# 软排序 (regularization_strength 越小，结果越接近真实排序)
sorted_x = torchsort.soft_sort(x, regularization_strength=0.1)
print(sorted_x)
# 输出近似为：tensor([[-0., 1., 2., 3., 5., 6., 7., 8., 9.]])

# 软排名
ranked_x = torchsort.soft_rank(x)
print(ranked_x)
# 输出：tensor([[8., 1., 5., 4., 3., 2., 6., 7., 9.]])

2. GPU 加速与反向传播

Torchsort 原生支持 GPU 和张量梯度计算，可直接用于模型训练。

# 启用 GPU 并要求梯度
x = torch.tensor([[8., 0., 5., 3., 2., 1., 6., 7., 9.]], requires_grad=True).cuda()

# 执行软排序
y = torchsort.soft_sort(x)

# 计算梯度
grad = torch.autograd.grad(y[0, 0], x)
print(grad)
# 输出梯度张量 (device='cuda:0')

3. 实际应用示例：可微分的 Spearman 相关系数

利用可微分排名特性，可以直接优化模型的排序指标。

def spearmanr(pred, target, **kw):
    # 获取软排名
    pred = torchsort.soft_rank(pred, **kw)
    target = torchsort.soft_rank(target, **kw)
    
    # 标准化处理
    pred = pred - pred.mean()
    pred = pred / pred.norm()
    target = target - target.mean()
    target = target / target.norm()
    
    return (pred * target).sum()

# 模拟预测值和目标值
pred = torch.tensor([[1., 2., 3., 4., 5.]], requires_grad=True)
target = torch.tensor([[5., 6., 7., 8., 7.]])

# 计算损失并反向传播
loss = 1.0 - spearmanr(pred, target)
loss.backward()