Google在去年初的论文中首次推出了时间复杂度为O(nlogn)的O(nlogn)算法，O(n)具有空间复杂度的可微排序算法，这比现有方法快一个数量级！最近有人在GitHub上开源了一个项目，以软件包的形式实现了快速的可微分排序和排名。上线几天，收获了300+星。项目地址：https://github.com/teddykoker/torchsort《Fast Differentiable Sorting and Ranking》论文地址：https://arxiv.org/pdf/2002.08871.pdfTorchsortTorchsort实现了Blondel等人提出的FastDifferentiableSortingandRanking。和Ranking），是基于纯PyTorch实现的。大部分代码是从项目“google-research/fast-soft-sort”中的原始Numpy实现复制而来的，带有自定义C++和CUDA内核以实现快速性能。Torchsort的安装方法很简单，使用常见的pip安装即可，安装代码如下：pipinstalltorchsort如果要构建CUDA扩展，需要安装CUDA工具链。如果要构建docker等没有CUDA运行环境的应用，安装前需要导出环境变量“TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing"”。使用方法torchsort有两个函数：soft_rank和soft_sort，每个函数都有参数regularization(l2或kl)（正则化函数）和regularization_strength（标量值）。每个都会对二维张量的最后一个维度进行排序，精度取决于正则化强度：importtorchiimporttorchsortx=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)[[-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创建一个可微分的Spearman等级系数函数，这样模型就可以直接针对这个指标进行优化：.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]]),)benchmarktorchsort和fast_soft_sort这两个操作的时间复杂度为O(nlogn)，相比内置的torch.sort，每个操作有一些额外的开销。NumbaJIT的批量大小为1（见左图）时，fast_soft_sort的前向传递与TorchsortCPU核心的性能大致相同，但它的反向传递仍然依赖于一些Python代码，这大大降低了它的性能。此外，torchsort内核支持批处理，随着批大小的增加，它会产生比fast_soft_sort更好的性能。torchsortCUDA内核在序列长度低于2000的情况下表现良好，并且可以扩展到非常大的批次。未来，CUDA内核可能会进一步优化，以达到接近内置torch.sort的性能。
