快速可微分排序算法PyTorch包，配有自定义C ++和CUDA，性能更好

有人将快速可微分排序算法打包实现，性能还不错。

谷歌去年年初在论文《Fast Differentiable Sorting and Ranking》中，重磅推出了首个具有 O(nlogn) 时间复杂度、O(n) 空间复杂度可微分排序算法，速度比现有方法快出一个数量级！

近日，有人在 GitHub 上开源了一个项目，通过软件包的形式实现了快速可微分排序和排名，上线几天，收获 300 + 星。

项目地址：https://github.com/teddykoker/torchsort
《Fast Differentiable Sorting and Ranking》论文地址：https://arxiv.org/pdf/2002.08871.pdf

Torchsort

Torchsort 实现了 Blondel 等人提出的快速可微分排序和排名（Fast Differentiable Sorting and Ranking），是基于纯 PyTorch 实现的。大部分代码是在项目「google-research/fast-soft-sort」中的原始 Numpy 实现复制而来，并配有自定义 C ++ 和 CUDA 内核以实现快速性能。

Torchsort 安装方式非常简单，采用常用的 pip 安装即可，安装代码如下：

pip install torchsort

如果你想构建 CUDA 扩展，你需要安装 CUDA 工具链。如果你想在没有 CUDA 运行环境中构建如 docker 的应用，在安装前需要导出环境变量「TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing"」。

使用方法

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'),)

示例展示

斯皮尔曼等级系数是用于测量两个变量之间单调相关性的非常有用的指标。我们可以使用 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 相比，每个操作都具有一些额外的开销。Numba JIT 的批处理大小为 1（请参见左图），fast_soft_sort 的前向传递与 Torchsort CPU 内核的性能大致相同，但是其后向传递仍然依赖于某些 Python 代码，这极大地降低了其性能。

此外，torchsort 内核支持批处理，随着批处理大小的增加，会产生比 fast_soft_sort 更好的性能。

torchsort CUDA 内核在序列长度低于 2000 时表现出色，并且可以扩展到非常大的 batch。在未来，CUDA 内核可能会进一步优化，以达到接近内置的 torch.sort 的性能。

工程GitHub可微分PyTorch

相关数据

排序算法技术

排序算法是将一串数据依照特定排序方式进行排列的算法，最常用到的排序方式是数值顺序以及字典顺序。基本上，排序算法的输出必须遵守下列两个原则：输出结果为递增序列（递增是针对所需的排序顺序而言）；输出结果是原输入的一种排列、或是重组。

来源：维基百科

基准技术

一种简单的模型或启发法，用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

来源：Google ML Glossary

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

时间复杂度技术

在计算机科学中，算法的时间复杂度是一个函数，它定量描述了该算法的运行时间。这是一个代表算法输入值的字符串的长度的函数。时间复杂度常用大O符号表述，不包括这个函数的低阶项和首项系数。使用这种方式时，时间复杂度可被称为是渐近的，亦即考察输入值大小趋近无穷时的情况。例如，如果一个算法对于任何大小为 n （必须比 n0 大）的输入，它至多需要 5n3 + 3n 的时间运行完毕，那么它的渐近时间复杂度是 O(n3)。

来源：维基百科

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数，这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在维空间内，有个分量的一种量，其中每个分量都是坐标的函数，而在坐标变换时，这些分量也依照某些规则作线性变换。称为该张量的秩或阶（与矩阵的秩和阶均无关系）。在数学里，张量是一种几何实体，或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达，记作标量的数组，但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中，表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了，它们都是二阶张量，对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

来源：维基百科

准确率技术

分类模型的正确预测所占的比例。在多类别分类中，准确率的定义为：正确的预测数/样本总数。在二元分类中，准确率的定义为：(真正例数+真负例数)/样本总数

来源：Google ML Glossary

正则化技术

当模型的复杂度增大时，训练误差会逐渐减小并趋向于0；而测试误差会先减小，达到最小值后又增大。当选择的模型复杂度过大时，过拟合现象就会发生。这样，在学习时就要防止过拟合。进行最优模型的选择，即选择复杂度适当的模型，以达到使测试误差最小的学习目的。

来源：李航著统计学习方法清华大学出版社