Fast, differentiable sorting and ranking in PyTorch

Overview

Torchsort

Tests

Fast, differentiable sorting and ranking in PyTorch.

Pure PyTorch implementation of Fast Differentiable Sorting and Ranking (Blondel et al.). Much of the code is copied from the original Numpy implementation at google-research/fast-soft-sort, with the isotonic regression solver rewritten as a PyTorch C++ and CUDA extension.

Install

pip install torchsort

To 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 variable TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing" before installing.

Usage

torchsort exposes two functions: soft_rank and soft_sort, each with parameters regularization ("l2" or "kl") and regularization_strength (a scalar value). Each will rank/sort the last dimension of a 2-d tensor, with an accuracy dependant upon the 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.]])

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

Example

Spearman's Rank Coefficient

Spearman's rank coefficient is 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:

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

Benchmark

Benchmark

torchsort and fast_soft_sort each operate with a time complexity of O(n log n), each with some additional overhead when compared to the built-in torch.sort. With a batch size of 1 (see left), the Numba JIT'd forward pass of fast_soft_sort performs about on-par with the torchsort CPU kernel, however its backward pass still relies on some Python code, which greatly penalizes its performance.

Furthermore, the torchsort kernel supports batches, and yields much better performance than fast_soft_sort as the batch size increases.

Benchmark

The torchsort CUDA 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 in torch.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}
}
Issues
  • pip install failed in windows

    pip install failed in windows

    Hi, I faced an installation error in windows. It installed fine in my ubuntu system. Could you tell me how I can fix it?

    Internal error: assertion failed at: "C:/dvs/p4/build/sw/rel/gpu_drv/r400/r400_00/drivers/compiler/edg/EDG_4.14/
    src/decl_spec.c", line 9596
        
        
        1 catastrophic error detected in the compilation of "C:/Users/Reasat/AppData/Local/Temp/tmpxft_000028b0_00000000
    -5_isotonic_cuda.cpp4.ii".
        Compilation aborted.
        isotonic_cuda.cu
        nvcc error   : 'cudafe++' died with status 0xC0000409
        error: command 'C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v10.0\\bin\\nvcc.exe' failed with exit st
    atus 9
        Error in atexit._run_exitfuncs:
        Traceback (most recent call last):
          File "C:\Users\Reasat\AppData\Roaming\Python\Python37\site-packages\colorama\ansitowin32.py", line 59, in clos
    ed
            return stream.closed
        ValueError: underlying buffer has been detached
        ----------------------------------------
    ERROR: Command errored out with exit status 1: 'D:\Miniconda3\envs\pytorch\python.exe' -u -c 'import sys, setuptools
    , tokenize; sys.argv[0] = '"'"'C:\\Users\\Reasat\\AppData\\Local\\Temp\\pip-install-uw3b5i5w\\torchsort_f8d66d1aaac6
    44a78d37585cc7273f94\\setup.py'"'"'; __file__='"'"'C:\\Users\\Reasat\\AppData\\Local\\Temp\\pip-install-uw3b5i5w\\to
    rchsort_f8d66d1aaac644a78d37585cc7273f94\\setup.py'"'"';f=getattr(tokenize, '"'"'open'"'"', open)(__file__);code=f.r
    ead().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' install --recor
    d 'C:\Users\Reasat\AppData\Local\Temp\pip-record-nbabfu8b\install-record.txt' --single-version-externally-managed --
    compile --install-headers 'D:\Miniconda3\envs\pytorch\Include\torchsort' Check the logs for full command output.
    
    opened by Reasat 12
  • cuda TypeError: 'NoneType' object is not callable

    cuda TypeError: 'NoneType' object is not callable

    >>> import torch
    >>> import torchsort
    >>> x = torch.tensor([[8., 0., 5., 3., 2., 1., 6., 7., 9.]], requires_grad=True).cuda()
    >>> y = torchsort.soft_sort(x)
    
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
      File "/home/shuiy/anaconda3/envs/pytorch_py3/lib/python3.7/site-packages/torchsort/ops.py", line 48, in soft_sort
        return SoftSort.apply(values, regularization, regularization_strength)
      File "/home/shuiy/anaconda3/envs/pytorch_py3/lib/python3.7/site-packages/torchsort/ops.py", line 132, in forward
        sol = isotonic_l2[s.device.type](w - s)
    TypeError: 'NoneType' object is not callable
    

    on jupyter notebook is:

    ---------------------------------------------------------------------------
    TypeError                                 Traceback (most recent call last)
    /tmp/ipykernel_8938/1075883647.py in <module>
          1 x = torch.tensor([[8., 0., 5., 3., 2., 1., 6., 7., 9.]], requires_grad=True).cuda()
    ----> 2 y = torchsort.soft_sort(x)
    
    ~/anaconda3/envs/pytorch_py3/lib/python3.7/site-packages/torchsort/ops.py in soft_sort(values, regularization, regularization_strength)
         46     if regularization not in ["l2", "kl"]:
         47         raise ValueError(f"'regularization' should be a 'l2' or 'kl'")
    ---> 48     return SoftSort.apply(values, regularization, regularization_strength)
         49 
         50 
    
    ~/anaconda3/envs/pytorch_py3/lib/python3.7/site-packages/torchsort/ops.py in forward(ctx, tensor, regularization, regularization_strength)
        130         # note reverse order of args
        131         if ctx.regularization == "l2":
    --> 132             sol = isotonic_l2[s.device.type](w - s)
        133         else:
        134             sol = isotonic_kl[s.device.type](w, s)
    
    TypeError: 'NoneType' object is not callable
    

    if x is on cpu(), run code is ok python 3.7.10, pytorch 1.9.0 , cudatoolkit=11.1, ubuntu 18.04

    opened by shuiyuejihua 11
  • Problem installing when no GPU present (in docker build step for example)

    Problem installing when no GPU present (in docker build step for example)

    Doesn't install during docker build phase (that does not have GPUs configured).

    Get error: /root/miniconda/lib/python3.8/site-packages/torch/cuda/init.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at /opt/conda/conda-bld/pytorch_1607370172916/work/c10/cuda/CUDAFunctions.cpp:100.) return torch._C._cuda_getDeviceCount() > 0

    If I install on the same image after running it with GPUs enabled it installs fine.

    opened by pcnudde 8
  • Unable to install with CUDA

    Unable to install with CUDA

    Hi, I'm excited to use this package but unfortunately am having issues getting it working with CUDA. I am using a conda env and have followed the steps in the README related to that. My torch version is 1.11.0 and my cudatoolkit version is 11.3.1. My Python version is 3.8.13 on a Linux machine if that is relevant.

    From a fresh environment:

    conda install -c pytorch pytorch torchvision cudatoolkit=11.3
    pip install torchsort 
    

    then if I try to use torchsort on a CUDA tensor, I get ImportError: You are trying to use the torchsort CUDA extension, but it looks like it is not available. Make sure you have the CUDA toolchain installed, and reinstall torchsort withpip install --force-reinstall --no-cache-dir torchsortto rebuild the extension. (which I have tried a few times now).

    Any help in getting this working would be amazing! Thanks so much!

    opened by sachit-menon 7
  • Hi, NVIDIA CUDA version >= 11.4 does not seem to install successfully.

    Hi, NVIDIA CUDA version >= 11.4 does not seem to install successfully.

    I tried to install the package in Tesla A100 and GeForce RTX 3090 with CUDA version 11.4 both failed. Can you provide some help please? Thank you very much!

    opened by XiaoqiWang 7
  • Any help a pip3 install --user issue?

    Any help a pip3 install --user issue?

    Here is my error code:

    Using legacy 'setup.py install' for torchsort, since package 'wheel' is not installed.
    Installing collected packages: torchsort
        Running setup.py install for torchsort ... \       error
    ERROR: Command errored out with exit status 1:  
         command: /share/software/user/open/python/3.9.0/bin/python3.9 -u -c 'import io, os, sys, setuptools, tokenize; sys.argv[0] = '"'"'/tmp/pip-install-rc2snp_l/torchsort_02ab43cb664b4f778f8057965fe69d12/setup.py'"'"'; __file__='"'"'/tmp/pip-install-rc2snp_l/torchsort_02ab43cb664b4f778f8057965fe69d12/setup.py'"'"';f = getattr(tokenize, '"'"'open'"'"', open)(__file__) if os.path.exists(__file__) else io.StringIO('"'"'from setuptools import setup; setup()'"'"');code = f.read().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' install --record /tmp/pip-record-0etipwtv/install-record.txt --single-version-externally-managed --user --prefix= --compile --install-headers /home/users/huangda/.local/include/python3.9/torchsort
             cwd: /tmp/pip-install-rc2snp_l/torchsort_02ab43cb664b4f778f8057965fe69d12/
        Complete output (31 lines):  
        No CUDA runtime is found, using CUDA_HOME='/share/software/user/open/cuda/11.2.0'  
        running install  
        running build  
        running build_py  
        creating build  
        creating build/lib.linux-x86_64-3.9  
        creating build/lib.linux-x86_64-3.9/torchsort  
        copying torchsort/__init__.py -> build/lib.linux-x86_64-3.9/torchsort  
        copying torchsort/ops.py -> build/lib.linux-x86_64-3.9/torchsort  
        running egg_info  
        writing torchsort.egg-info/PKG-INFO  
        writing dependency_links to torchsort.egg-info/dependency_links.txt    
        writing requirements to torchsort.egg-info/requires.txt  
        writing top-level names to torchsort.egg-info/top_level.txt  
        /share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/utils/cpp_extension.py:369: UserWarning: Attempted to use ninja as the BuildExtension backend but we could not find ninja.. Falling back to using the slow distutils backend.  
          warnings.warn(msg.format('we could not find ninja.'))  
        reading manifest file 'torchsort.egg-info/SOURCES.txt'  
        reading manifest template 'MANIFEST.in'  
        writing manifest file 'torchsort.egg-info/SOURCES.txt'  
        copying torchsort/isotonic_cpu.cpp -> build/lib.linux-x86_64-3.9/torchsort  
        copying torchsort/isotonic_cuda.cu -> build/lib.linux-x86_64-3.9/torchsort  
        running build_ext   
        building 'torchsort.isotonic_cpu' extension  
        creating build/temp.linux-x86_64-3.9  
        creating build/temp.linux-x86_64-3.9/torchsort  
        gcc -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -fPIC -I/share/software/user/open/py- 
     pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/include -I/share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -I/share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/include/TH -I/share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/include/THC -I/share/software/user/open/python/3.9.0/include/python3.9 -c torchsort/isotonic_cpu.cpp -o build/temp.linux-x86_64-3.9/torchsort/isotonic_cpu.o -fopenmp -ffast-math -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE="_gcc" -DPYBIND11_STDLIB="_libstdcpp" -DPYBIND11_BUILD_ABI="_cxxabi1011" -DTORCH_EXTENSION_NAME=isotonic_cpu -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14
        c++ -pthread -shared -L/share/software/user/open/libffi/3.2.1/lib64 -L/share/software/user/open/libressl/3.2.1/lib -L/share/software/user/open/sqlite/3.18.0/lib -L/share/software/user/open/tcltk/8.6.6/lib -L/share/software/user/open/xz/5.2.3/lib -L/share/software/user/open/zlib/1.2.11/lib build/temp.linux-x86_64-3.9/torchsort/isotonic_cpu.o -L/share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/lib -L/share/software/user/open/python/3.9.0/lib -lc10 -ltorch -ltorch_cpu -ltorch_python -o build/lib.linux-x86_64-3.9/torchsort/isotonic_cpu.cpython-39-x86_64-linux-gnu.so
        building 'torchsort.isotonic_cuda' extension
        /share/software/user/open/cuda/11.2.0/bin/nvcc -I/share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/include -I/share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -I/share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/include/TH -I/share/software/user/open/py-pytorch/1.8.1_py39/lib/python3.9/site-packages/torch/include/THC -I/share/software/user/open/cuda/11.2.0/include -I/share/software/user/open/python/3.9.0/include/python3.9 -c torchsort/isotonic_cuda.cu -o build/temp.linux-x86_64-3.9/torchsort/isotonic_cuda.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE="_gcc" -DPYBIND11_STDLIB="_libstdcpp" -DPYBIND11_BUILD_ABI="_cxxabi1011" -DTORCH_EXTENSION_NAME=isotonic_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_61,code=compute_61 -gencode=arch=compute_61,code=sm_61 -gencode=arch=compute_70,code=compute_70 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=compute_75 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 -ccbin gcc -std=c++14
        nvcc error   : 'cicc' died due to signal 9 (Kill signal)  
        error: command '/share/software/user/open/cuda/11.2.0/bin/nvcc' failed with exit code 9  
        ----------------------------------------
    ERROR: Command errored out with exit status 1: /share/software/user/open/python/3.9.0/bin/python3.9 -u -c 'import io, os, sys, setuptools, tokenize; sys.argv[0] = '"'"'/tmp/pip-install-rc2snp_l/torchsort_02ab43cb664b4f778f8057965fe69d12/setup.py'"'"'; __file__='"'"'/tmp/pip-install-rc2snp_l/torchsort_02ab43cb664b4f778f8057965fe69d12/setup.py'"'"';f = getattr(tokenize, '"'"'open'"'"', open)(__file__) if os.path.exists(__file__) else io.StringIO('"'"'from setuptools import setup; setup()'"'"');code = f.read().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' install --record /tmp/pip-record-0etipwtv/install-record.txt --single-version-externally-managed --user --prefix= --compile --install-headers /home/users/huangda/.local/include/python3.9/torchsort Check the logs for full command output.
    

    Has anybody else experienced this before? The full command I am running is:

    TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing;Ampere" pip3 install --user torchsort

    opened by derekahuang 6
  • soft_rank has memory leak?

    soft_rank has memory leak?

    Hi I have installed the main branch, and I'm seeing that the torchsort.soft_rank function is causing memory leaks. Looking at, nvidia-smi, it does not free up any memory and I see the following printed out over and over:

    [W pthreadpool-cpp.cc:90] Warning: Leaking Caffe2 thread-pool after fork. (function pthreadpool)
    
    opened by ashkan-leo 6
  • RuntimeError: CUDA error: an illegal memory access was encountered

    RuntimeError: CUDA error: an illegal memory access was encountered

    I got the below error when training my network for a while (10-20 epochs). Traceback (most recent call last): File "train.py", line 232, in <module> main() File "train.py", line 179, in main train(cfg, train_loader, model, criterion, optimizer, lr_scheduler, epoch, final_output_dir, tb_log_dir, writer_dict) File "/home/maxchu/Fin/numerai_dev/function.py", line 62, in train loss, loss_indv = criterion(pred, target, auto_pred, auto_target) File "/home/maxchu/fin_venv/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl result = self.forward(*input, **kwargs) File "/home/maxchu/Fin/numerai_dev/loss.py", line 39, in forward - 0.1 * spearman(pred, target, regularization_strength=1e-2), File "/home/maxchu/Fin/numerai_dev/loss.py", line 24, in spearman pred = torchsort.soft_rank( File "/home/maxchu/fin_venv/lib/python3.8/site-packages/torchsort-0.1.4-py3.8-linux-x86_64.egg/torchsort/ops.py", line 40, in soft_rank return SoftRank.apply(values, regularization, regularization_strength) File "/home/maxchu/fin_venv/lib/python3.8/site-packages/torchsort-0.1.4-py3.8-linux-x86_64.egg/torchsort/ops.py", line 96, in forward ret = (s - dual_sol).gather(1, inv_permutation) RuntimeError: CUDA error: an illegal memory access was encountered Some details:

    1. System: Ubuntu 18.06
    2. Python 3.8 using venv
    3. Install method: Manual compile (git clone -> python setup.py install)
    4. PyTorch 1.8.0 with cuda 10.2 (and correspinding pytorch geometric package)

    Please let me know if you need more informations.

    opened by MaxChu719 6
  • isotonic_cpu.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZNK2at6Tensor8data_ptrIfEEPT_v

    isotonic_cpu.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZNK2at6Tensor8data_ptrIfEEPT_v

    Hello, thank you for the library. I am trying to use spearmanr example from the main README.md on V100s and A100s GPUs but getting error below. cuda 11.3 pytorch 1.10.0 python 3.7.11 torchsort 0.1.7

    File "/truba/home/fkahraman/miniconda3/envs/openmmlab/lib/python3.7/site-packages/torchsort/ops.py", line 18, in from .isotonic_cpu import isotonic_kl as isotonic_kl_cpu ImportError: /truba/home/fkahraman/miniconda3/envs/openmmlab/lib/python3.7/site-packages/torchsort/isotonic_cpu.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZNK2at6Tensor8data_ptrIfEEPT_v from .isotonic_cpu import isotonic_kl as isotonic_kl_cpu ImportError: /truba/home/fkahraman/miniconda3/envs/openmmlab/lib/python3.7/site-packages/torchsort/isotonic_cpu.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZNK2at6Tensor8data_ptrIfEEPT_v ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 99273) of binary: /truba/home/fkahraman/miniconda3/envs/openmmlab/bin/python

    opened by fehmikahraman 5
  • Failed building wheel when trying to install

    Failed building wheel when trying to install

    On my potato GPU-less desktop pip install torchsort works great in the anaconda prompt, but on another machine I get a "failed building wheel" followed by a huge error message. The environment works fine with lots of libraries and I'm running a project with it, it's just torchsort that does this problem for me. Do you have any clue what could cause that? Thanks!

    opened by roeeben 5
  • Incorrect results when running on non-default cuda device

    Incorrect results when running on non-default cuda device

    When running torchsort.soft_rank or torchsort.soft_sort on a tensor that's not on the default cuda device (usually cuda:0), the results are incorrect.

    import torch
    import torchsort
    
    x = torch.tensor([[9,8]], device="cuda:1")
    
    print(torchsort.soft_rank(x))
    # tensor([[9., 8.]], device='cuda:1')
    
    print(torchsort.soft_sort(x))
    # tensor([[-2., -1.]], device='cuda:1')
    

    Based on the GPU memory usage, torchsort tries to do something on the default cuda device cuda:0 instead of whichever device the input tensor is on. As a workaround, you need to either change the default cuda device with torch.cuda.set_device or use the context manager torch.cuda.device.

    import torch
    import torchsort
    
    x = torch.tensor([[9,8]], device="cuda:1")
    
    with torch.cuda.device(x.device):
        print(torchsort.soft_rank(x))
        # tensor([[2., 1.]], device='cuda:1')
    
        print(torchsort.soft_sort(x))
        # tensor([[8., 9.]], device='cuda:1')
    
    opened by JustinSzeto 4
  • CUDA benchmarks might be misleading

    CUDA benchmarks might be misleading

    I wanted to try to improve/modify the torchsort code a little so I tried making a copy of the SoftSort class and the soft_sort function.

    Running some benchmarks I got the following results: benchmark_custom benchmark_custom_cuda

    Which was worrying. The carbon copy diverges at a similar point to the figure in the readme:

    I then re-ran the benchmark with the exact same function twice (not even a copy) and got the same results.

    That code can be found here:

    import sys
    from collections import defaultdict
    from timeit import timeit
    
    import matplotlib.pyplot as plt
    import torch
    
    import torchsort
    
    try:
        import fast_soft_sort.pytorch_ops as fss
    except ImportError:
        print("install fast_soft_sort:")
        print("pip install git+https://github.com/google-research/fast-soft-sort")
        sys.exit()
    
    
    N = list(range(1, 5_000, 100))
    B = [2 ** i for i in range(9)]
    B_CUDA = [2 ** i for i in range(13)]
    SAMPLES = 100
    CONVERT = 1e-6  # convert seconds to micro-seconds
    
    
    def time(f):
        return timeit(f, number=SAMPLES) / SAMPLES / CONVERT
    
    
    def backward(f, x):
        y = f(x)
        torch.autograd.grad(y.sum(), x)
    
    
    def style(name):
        if name == "torch.sort":
            return {"color": "blue"}
        linestyle = "--" if "backward" in name else "-"
        if "fast_soft_sort" in name:
            return {"color": "green", "linestyle": linestyle}
        elif "again" in name:
            return {"color": "red", "linestyle": linestyle}
        else:
            return {"color": "orange", "linestyle": linestyle}
    
    
    def batch_size(ax):
        data = defaultdict(list)
        for b in B:
            x = torch.randn(b, 100)
            # data["torch.sort"].append(time(lambda: torch.sort(x)))
            data["torchsort"].append(time(lambda: torchsort.soft_sort(x)))
            data["torchsort_again"].append(time(lambda: torchsort.soft_sort(x)))
            # data["fast_soft_sort"].append(time(lambda: fss.soft_sort(x)))
            x = torch.randn(b, 100, requires_grad=True)
            data["torchsort (with backward)"].append(
                time(lambda: backward(torchsort.soft_sort, x))
            )
            data["torchsort_again (with backward)"].append(
                time(lambda: backward(torchsort.soft_sort, x))
            )
            # data["fast_soft_sort (with backward)"].append(
            #     time(lambda: backward(fss.soft_sort, x))
            # )
    
        for label in data.keys():
            ax.plot(B, data[label], label=label, **style(label))
        ax.set_xlabel("Batch Size")
        ax.set_ylim(0, 5000)
        ax.set_ylabel("Execution Time (μs)")
        ax.legend()
    
    
    def sequence_length(ax):
        data = defaultdict(list)
        for n in N:
            x = torch.randn(1, n)
            # data["torch.sort"].append(time(lambda: torch.sort(x)))
            data["torchsort"].append(time(lambda: torchsort.soft_sort(x)))
            data["torchsort_again"].append(time(lambda: torchsort.soft_sort(x)))
            # data["fast_soft_sort"].append(time(lambda: fss.soft_sort(x)))
            x = torch.randn(1, n, requires_grad=True)
            data["torchsort (with backward)"].append(
                time(lambda: backward(torchsort.soft_sort, x))
            )
            data["torchsort_again (with backward)"].append(
                time(lambda: backward(torchsort.soft_sort, x))
            )
            # data["fast_soft_sort (with backward)"].append(
            #     time(lambda: backward(fss.soft_sort, x))
            # )
    
        for label in data.keys():
            ax.plot(N, data[label], label=label, **style(label))
        ax.set_xlabel("Sequence Length")
        ax.set_ylim(0, 1000)
        ax.set_ylabel("Execution Time (μs)")
        ax.legend()
    
    
    def batch_size_cuda(ax):
        data = defaultdict(list)
        for b in B_CUDA:
            x = torch.randn(b, 100).cuda()
            # data["torch.sort"].append(time(lambda: torch.sort(x)))
            data["torchsort"].append(time(lambda: torchsort.soft_sort(x)))
            data["torchsort_again"].append(time(lambda: torchsort.soft_sort(x)))
            x = torch.randn(b, 100, requires_grad=True).cuda()
            data["torchsort (with backward)"].append(
                time(lambda: backward(torchsort.soft_sort, x))
            )
            data["torchsort_again (with backward)"].append(
                time(lambda: backward(torchsort.soft_sort, x))
            )
        for label in data.keys():
            ax.plot(B_CUDA, data[label], label=label, **style(label))
        ax.set_xlabel("Batch Size")
        ax.set_ylabel("Execution Time (μs)")
        ax.legend()
    
    
    def sequence_length_cuda(ax):
        data = defaultdict(list)
        for n in N:
            x = torch.randn(1, n).cuda()
            # data["torch.sort"].append(time(lambda: torch.sort(x)))
            data["torchsort"].append(time(lambda: torchsort.soft_sort(x)))
            data["torchsort_again"].append(time(lambda: torchsort.soft_sort(x)))
            x = torch.randn(1, n, requires_grad=True).cuda()
            data["torchsort (with backward)"].append(
                time(lambda: backward(torchsort.soft_sort, x))
            )
            data["torchsort_again (with backward)"].append(
                time(lambda: backward(torchsort.soft_sort, x))
            )
        for label in data.keys():
            ax.plot(N, data[label], label=label, **style(label))
        ax.set_xlabel("Sequence Length")
        ax.set_ylabel("Execution Time (μs)")
        ax.legend()
    
    
    if __name__ == "__main__":
        # jit/warmup
        x = torch.randn(1, 10, requires_grad=True)
        backward(torchsort.soft_sort, x)
        backward(fss.soft_sort, x)
    
        fig, (ax1, ax2) = plt.subplots(figsize=(10, 4), ncols=2)
        sequence_length(ax1)
        batch_size(ax2)
        fig.suptitle("Torchsort Benchmark: CPU")
        fig.tight_layout()
        plt.savefig("extra/benchmark3.png")
    
        if torch.cuda.is_available():
            # warmup
            x = torch.randn(1, 10, requires_grad=True).cuda()
            backward(torchsort.soft_sort, x)
    
            fig, (ax1, ax2) = plt.subplots(figsize=(10, 4), ncols=2)
            sequence_length_cuda(ax1)
            batch_size_cuda(ax2)
            fig.suptitle("Torchsort Benchmark: CUDA")
            fig.tight_layout()
            plt.savefig("extra/benchmark_cuda3.png")
    

    Any idea what this might depend on?

    opened by zimonitrome 1
  • Can I sort by specific column?

    Can I sort by specific column?

    Is there any way to sort a tensor by a given column?

    For example, soring by first column:

    input_tensor = torch.tensor([
            [1, 5], 
            [30, 30], 
            [6, 9], 
            [80, -2]
    ])
    
    target_tensor = torch.tensor([
            [80, -2],
            [30, 30], 
            [6, 9], 
            [1, 5], 
    ])
    
    enhancement 
    opened by zimonitrome 7
  • Reproducing CIFAR results

    Reproducing CIFAR results

    Thanks a lot for this implementation. I was wondering how can I use the repo to reproduce the results on CIFAR as reported in the paper. As I understand, the target one-hot encoding will serve as top-k classification(k=1). But, after obtaining the logits and passing through the softmax(putting output [0, 1]) the objective is to make the output follow the target ordering. How can this be achieved?

    question 
    opened by paganpasta 8
Releases(v0.1.9)
Owner
Teddy Koker
Machine Learning Researcher
Teddy Koker
A fast, scalable, high performance Gradient Boosting on Decision Trees library, used for ranking, classification, regression and other machine learning tasks for Python, R, Java, C++. Supports computation on CPU and GPU.

Website | Documentation | Tutorials | Installation | Release Notes CatBoost is a machine learning method based on gradient boosting over decision tree

CatBoost 6.6k Jun 21, 2022
A fast, distributed, high performance gradient boosting (GBT, GBDT, GBRT, GBM or MART) framework based on decision tree algorithms, used for ranking, classification and many other machine learning tasks.

Light Gradient Boosting Machine LightGBM is a gradient boosting framework that uses tree based learning algorithms. It is designed to be distributed a

Microsoft 13.9k Jun 18, 2022
Learning embeddings for classification, retrieval and ranking.

StarSpace StarSpace is a general-purpose neural model for efficient learning of entity embeddings for solving a wide variety of problems: Learning wor

Facebook Research 3.8k Jun 15, 2022
ADOP: Approximate Differentiable One-Pixel Point Rendering

ADOP: Approximate Differentiable One-Pixel Point Rendering

Darius Rückert 1.8k Jun 24, 2022
Efficient Differentiable Simulation of Articulated Bodies (ICML2021)

Efficient Differentiable Simulation of Articulated Bodies Yi-Ling Qiao, Junbang Liang, Vladlen Koltun, Ming C. Lin [Paper] [Video] [Slides] [Code] Set

YilingQiao 70 Jun 17, 2022
Deep Learning API and Server in C++11 support for Caffe, Caffe2, PyTorch,TensorRT, Dlib, NCNN, Tensorflow, XGBoost and TSNE

Open Source Deep Learning Server & API DeepDetect (https://www.deepdetect.com/) is a machine learning API and server written in C++11. It makes state

JoliBrain 2.4k Jun 13, 2022
PSTensor provides a way to hack the memory management of tensors in TensorFlow and PyTorch by defining your own C++ Tensor Class.

PSTensor : Custimized a Tensor Data Structure Compatible with PyTorch and TensorFlow. You may need this software in the following cases. Manage memory

Jiarui Fang 8 Feb 12, 2022
This repository is a tensorrt deployment of the onsets and frames model, which is implemented using pytorch.

Onsets and Frames TensorRT inference This repository is a tensorrt deployment of the onsets and frames model, which is implemented using pytorch (http

Xianke Wang 6 Jan 13, 2022
A c++ trainable semantic segmentation library based on libtorch (pytorch c++). Backbone: ResNet, ResNext. Architecture: FPN, U-Net, PAN, LinkNet, PSPNet, DeepLab-V3, DeepLab-V3+ by now.

中文 C++ library with Neural Networks for Image Segmentation based on LibTorch. The main features of this library are: High level API (just a line to cr

null 259 Jun 25, 2022
This is a code repository for pytorch c++ (or libtorch) tutorial.

LibtorchTutorials English version 环境 win10 visual sutdio 2017 或者Qt4.11.0 Libtorch 1.7 Opencv4.5 配置 libtorch+Visual Studio和libtorch+QT分别记录libtorch在VS和Q

null 323 Jun 27, 2022
GPU PyTorch TOP in TouchDesigner with CUDA-enabled OpenCV

PyTorchTOP This project demonstrates how to use OpenCV with CUDA modules and PyTorch/LibTorch in a TouchDesigner Custom Operator. Building this projec

David 65 Jun 15, 2022
functorch is a prototype of JAX-like composable function transforms for PyTorch.

functorch Why functorch? | Install guide | Transformations | Future Plans functorch is a prototype of JAX-like composable FUNCtion transforms for pyTO

Richard Zou 1k Jun 24, 2022
Support Yolov4/Yolov3/Centernet/Classify/Unet. use darknet/libtorch/pytorch to onnx to tensorrt

ONNX-TensorRT Yolov4/Yolov3/CenterNet/Classify/Unet Implementation Yolov4/Yolov3 centernet INTRODUCTION you have the trained model file from the darkn

null 156 Jun 10, 2022
functorch is a prototype of JAX-like composable function transforms for PyTorch.

functorch Why functorch? | Install guide | Transformations | Future Plans functorch is a prototype of JAX-like composable FUNCtion transforms for pyTO

Facebook Research 1k Jun 24, 2022
UE4 Plugin to execute trained PyTorch modules

SimplePyTorch UE4 Plugin to execute trained PyTorch modules ------- Packaging ------- Download PyTorch C++ distributions: https://pytorch.org/cppdocs/

null 41 Jun 28, 2022
C++ trainable detection library based on libtorch (or pytorch c++). Yolov4 tiny provided now.

C++ Library with Neural Networks for Object Detection Based on LibTorch. ?? Libtorch Tutorials ?? Visit Libtorch Tutorials Project if you want to know

null 44 Jun 23, 2022
A simple demonstration of how PyTorch autograd works

简单地演示了 PyTorch 中自动求导机制的原理。 官方博客:https://pytorch.org/blog/overview-of-pytorch-autograd-engine/ 编译运行 使用 Bazel bazel run autograd_test 包含了一个使用 MSE 损失函数的一

Howard Lau 14 Feb 24, 2022
An inofficial PyTorch implementation of PREDATOR based on KPConv.

PREDATOR: Registration of 3D Point Clouds with Low Overlap An inofficial PyTorch implementation of PREDATOR based on KPConv. The code has been tested

ZhuLifa 15 Mar 29, 2022
DLPrimitives/OpenCL out of tree backend for pytorch

Pytorch OpenCL backend based on dlprimitives DLPrimitives-OpenCL out of tree backend for pytorch It is only beginning, but you can train some vision n

Artyom Beilis 59 Jun 15, 2022