Package: torchsort
TorchsortFast, differentiable sorting and ranking in PyTorch.Pure PyTorch implementation ofFast Differentiable Sorting and
Ranking(Blondel et al.). Much of the code is
copied from the original Numpy implementation atgoogle-research/fast-soft-sort,
with the isotonic regression solver rewritten as a PyTorch C++ and CUDA
extension.InstallpipinstalltorchsortTo build the CUDA extension you will need the CUDA toolchain installed. If you
want to build in an environment without a CUDA runtime (e.g. docker), you will
need to export the environment variableTORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing;Ampere"before installing.Conda InstallationOn some systems the package my not compile with `pip` install in conda
environments. If this happens you may need to:Install g++ withconda install -c conda-forge gxx_linux-64Set export variableexport CXX=/path/to/miniconda3/envs/env_name/bin/x86_64-conda_cos6-linux-gnu-g++If still failing, export variableexport LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/miniconda3/libThanks to @levnikmyskin for pointing this out!Usagetorchsortexposes two functions:soft_rankandsoft_sort, each with
parametersregularization("l2"or"kl") andregularization_strength(a
scalar value). Each will rank/sort the last dimension of a 2-d tensor, with an
accuracy dependant upon the regularization strength:importtorchimporttorchsortx=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.]])Both operations are fully differentiable, on CPU or 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'),)ExampleSpearman's Rank CoefficientSpearman's rank
coefficientis a very useful metric for measuring how monotonically related two variables
are. We can use Torchsort to create a differentiable Spearman's rank coefficient
function so that we can optimize a model directly for this metric:importtorchimporttorchsortdefspearmanr(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]]),)Benchmarktorchsortandfast_soft_sorteach operate with a time complexity ofO(n log
n), each with some additional overhead when compared to the built-intorch.sort. With a batch size of 1 (see left), the Numba JIT'd forward pass offast_soft_sortperforms about on-par with thetorchsortCPU kernel, however
its backward pass still relies on some Python code, which greatly penalizes its
performance.Furthermore, thetorchsortkernel supports batches, and yields much better
performance thanfast_soft_sortas the batch size increases.ThetorchsortCUDA kernel performs quite well with sequence lengths under
~2000, and scales to extremely large batch sizes. In the future the
CUDA kernel can likely be further optimized to achieve performance closer to that of the
built intorch.sort.Reference@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}}