├── .coveragerc ├── .gitignore ├── .travis.yml ├── LICENSE ├── MANIFEST.in ├── README.md ├── cuda ├── unique.cpp └── unique_kernel.cu ├── setup.cfg ├── setup.py ├── test ├── __init__.py ├── test_unique.py └── utils.py └── torch_unique ├── __init__.py └── unique.py /.coveragerc: -------------------------------------------------------------------------------- 1 | [report] 2 | exclude_lines = 3 | pragma: no cover 4 | cuda 5 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | build/ 3 | dist/ 4 | .cache/ 5 | .eggs/ 6 | *.egg-info/ 7 | .coverage 8 | *.so 9 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | sudo: required 3 | dist: trusty 4 | matrix: 5 | include: 6 | - python: 2.7 7 | - python: 3.5 8 | - python: 3.6 9 | install: 10 | - if [[ $TRAVIS_PYTHON_VERSION == 2.7 ]]; then pip install http://download.pytorch.org/whl/cpu/torch-0.4.0-cp27-cp27mu-linux_x86_64.whl; fi 11 | - if [[ $TRAVIS_PYTHON_VERSION == 3.5 ]]; then pip install http://download.pytorch.org/whl/cpu/torch-0.4.0-cp35-cp35m-linux_x86_64.whl; fi 12 | - if [[ $TRAVIS_PYTHON_VERSION == 3.6 ]]; then pip install http://download.pytorch.org/whl/cpu/torch-0.4.0-cp36-cp36m-linux_x86_64.whl; fi 13 | - pip install pycodestyle 14 | - pip install flake8 15 | - pip install codecov 16 | script: 17 | - pycodestyle . 18 | - flake8 . 19 | - python setup.py install 20 | - python setup.py test 21 | after_success: 22 | - codecov 23 | notifications: 24 | email: false 25 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2018 Matthias Fey 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in 11 | all copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | THE SOFTWARE. 20 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include LICENSE 2 | recursive-include cuda * 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [pypi-image]: https://badge.fury.io/py/torch-unique.svg 2 | [pypi-url]: https://pypi.python.org/pypi/torch-unique 3 | [build-image]: https://travis-ci.org/rusty1s/pytorch_unique.svg?branch=master 4 | [build-url]: https://travis-ci.org/rusty1s/pytorch_unique 5 | [coverage-image]: https://codecov.io/gh/rusty1s/pytorch_unique/branch/master/graph/badge.svg 6 | [coverage-url]: https://codecov.io/github/rusty1s/pytorch_unique?branch=master 7 | 8 | # PyTorch Unique 9 | 10 | [![PyPI Version][pypi-image]][pypi-url] 11 | [![Build Status][build-image]][build-url] 12 | [![Code Coverage][coverage-image]][coverage-url] 13 | 14 | -------------------------------------------------------------------------------- 15 | 16 | **PyTorch 0.4.1 now supports [`unique`](https://pytorch.org/docs/stable/torch.html#torch.unique) both for CPU and GPU. 17 | Therefore, this package is no longer needed and will not be updated. 18 | In contrast to this package, PyTorch's version does not return an index array. 19 | However, you can easily generate it by using the following code:** 20 | 21 | ```python 22 | import torch 23 | 24 | unique, inverse = torch.unique(x, sorted=True, return_inverse=True) 25 | perm = torch.arange(inverse.size(0), dtype=inverse.dtype, device=inverse.device) 26 | inverse, perm = inverse.flip([0]), perm.flip([0]) 27 | perm = inverse.new_empty(unique.size(0)).scatter_(0, inverse, perm) 28 | ``` 29 | 30 | -------------------------------------------------------------------------------- 31 | 32 | This package consists of a small extension library of a highly optimized `unique` operation for the use in [PyTorch](http://pytorch.org/), which is missing in the main package. 33 | The operation works on varying data types and is implemented both for CPU and GPU. 34 | 35 | ## Installation 36 | 37 | Ensure that PyTorch 0.4.0 is installed and verify that `cuda/bin` and `cuda/include` are in your `$PATH` and `$CPATH` respectively, *e.g.*: 38 | 39 | ``` 40 | $ python -c "import torch; print(torch.__version__)" 41 | >>> 0.4.0 42 | 43 | $ echo $PATH 44 | >>> /usr/local/cuda/bin:... 45 | 46 | $ echo $CPATH 47 | >>> /usr/local/cuda/include:... 48 | ``` 49 | 50 | Then run: 51 | 52 | ``` 53 | pip install torch-unique 54 | ``` 55 | 56 | If you are running into any installation problems, please create an [issue](https://github.com/rusty1s/pytorch_unique/issues). 57 | Be sure to import `torch` first before using this package to resolve symbols the dynamic linker must see. 58 | 59 | ## Usage 60 | 61 | ``` 62 | torch_unique.unique(src) -> (Tensor, LongTensor) 63 | ``` 64 | 65 | Returns the sorted unique scalar elements of the input tensor as an one-dimensional tensor. 66 | 67 | A tuple of `(unique_tensor, unique_indices)` is returned, where the `unique_indices` are the indices of the elements in the original input tensor. Note that `unique_indices` is not guaranteed to be stable on GPU. 68 | 69 | ### Parameters 70 | 71 | * **src** *(Tensor)* - The input tensor. 72 | 73 | ### Returns 74 | 75 | * **out** *(Tensor)* - The unique elements from `src` as an one-dimensional tensor. 76 | * **perm** *(LongTensor)* - The unique indices from `src` as an one-dimensional tensor. 77 | 78 | ### Example 79 | 80 | ```py 81 | import torch 82 | from torch_unique import unique 83 | 84 | src = torch.tensor([100, 10, 100, 1, 1000, 1, 1000, 1]) 85 | out, perm = unique(src) 86 | ``` 87 | 88 | ``` 89 | print(out) 90 | tensor([ 1, 10, 100, 1000]) 91 | print(perm) 92 | tensor([ 3, 1, 0, 4]) 93 | ``` 94 | 95 | ## Running tests 96 | 97 | ``` 98 | python setup.py test 99 | ``` 100 | -------------------------------------------------------------------------------- /cuda/unique.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor") 4 | 5 | std::tuple unique_cuda(at::Tensor src); 6 | 7 | std::tuple unique(at::Tensor src) { 8 | CHECK_CUDA(src); 9 | return unique_cuda(src); 10 | } 11 | 12 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 13 | m.def("unique", &unique, "Unique (CUDA)"); 14 | } 15 | -------------------------------------------------------------------------------- /cuda/unique_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #define THREADS 1024 4 | #define BLOCKS(N) (N + THREADS - 1) / THREADS 5 | 6 | template 7 | __global__ void unique_cuda_kernel(scalar_t *__restrict__ src, uint8_t *mask, 8 | size_t numel) { 9 | const size_t index = blockIdx.x * blockDim.x + threadIdx.x; 10 | const size_t stride = blockDim.x * gridDim.x; 11 | for (ptrdiff_t i = index; i < numel; i += stride) { 12 | if (i == 0 || src[i] != src[i - 1]) { 13 | mask[i] = 1; 14 | } 15 | } 16 | } 17 | 18 | std::tuple unique_cuda(at::Tensor src) { 19 | at::Tensor perm; 20 | std::tie(src, perm) = src.sort(); 21 | 22 | auto mask = at::zeros(src.type().toScalarType(at::kByte), {src.numel()}); 23 | AT_DISPATCH_ALL_TYPES(src.type(), "grid_cuda_kernel", [&] { 24 | unique_cuda_kernel<<>>( 25 | src.data(), mask.data(), src.numel()); 26 | }); 27 | 28 | src = src.masked_select(mask); 29 | perm = perm.masked_select(mask); 30 | 31 | return std::make_tuple(src, perm); 32 | } 33 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | description-file = README.md 3 | 4 | [aliases] 5 | test = pytest 6 | 7 | [tool:pytest] 8 | addopts = --capture=no --cov 9 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME 3 | 4 | __version__ = '1.0.3' 5 | url = 'https://github.com/rusty1s/pytorch_unique' 6 | 7 | install_requires = ['numpy'] 8 | setup_requires = ['pytest-runner'] 9 | tests_require = ['pytest', 'pytest-cov'] 10 | ext_modules = [] 11 | cmdclass = {} 12 | 13 | if CUDA_HOME: 14 | ext_modules += [ 15 | CUDAExtension('torch_unique.unique_cuda', 16 | ['cuda/unique.cpp', 'cuda/unique_kernel.cu']) 17 | ] 18 | cmdclass['build_ext'] = BuildExtension 19 | 20 | setup( 21 | name='torch_unique', 22 | version=__version__, 23 | description='Optimized PyTorch Unique Operation', 24 | author='Matthias Fey', 25 | author_email='matthias.fey@tu-dortmund.de', 26 | url=url, 27 | download_url='{}/archive/{}.tar.gz'.format(url, __version__), 28 | keywords=['pytorch', 'unique', 'python'], 29 | install_requires=install_requires, 30 | setup_requires=setup_requires, 31 | tests_require=tests_require, 32 | ext_modules=ext_modules, 33 | cmdclass=cmdclass, 34 | packages=find_packages(), ) 35 | -------------------------------------------------------------------------------- /test/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rusty1s/pytorch_unique/228b5d245836934663b9b1d0264db18fe927be7d/test/__init__.py -------------------------------------------------------------------------------- /test/test_unique.py: -------------------------------------------------------------------------------- 1 | from itertools import product 2 | 3 | import pytest 4 | from torch_unique import unique 5 | 6 | from .utils import dtypes, devices, tensor 7 | 8 | 9 | @pytest.mark.parametrize('dtype,device', product(dtypes, devices)) 10 | def test_unique(dtype, device): 11 | src = tensor([100, 10, 100, 1, 200, 1, 200, 10], dtype, device) 12 | 13 | out, perm = unique(src) 14 | assert out.tolist() == [1, 10, 100, 200] 15 | assert out.tolist() == src[perm].tolist() 16 | -------------------------------------------------------------------------------- /test/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.testing import get_all_dtypes 3 | 4 | dtypes = get_all_dtypes() 5 | dtypes.remove(torch.half) 6 | dtypes.remove(torch.int8) # NumPy conversion for CharTensor is not supported 7 | 8 | devices = [torch.device('cpu')] 9 | if torch.cuda.is_available(): 10 | devices += [torch.device('cuda:{}'.format(torch.cuda.current_device()))] 11 | 12 | 13 | def tensor(x, dtype, device): 14 | return None if x is None else torch.tensor(x, dtype=dtype, device=device) 15 | -------------------------------------------------------------------------------- /torch_unique/__init__.py: -------------------------------------------------------------------------------- 1 | from .unique import unique 2 | 3 | __version__ = '1.0.3' 4 | 5 | __all__ = ['unique', '__version__'] 6 | -------------------------------------------------------------------------------- /torch_unique/unique.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | if torch.cuda.is_available(): 5 | import torch_unique.unique_cuda 6 | 7 | 8 | def unique(src): 9 | """Returns the sorted unique scalar elements of the input tensor as an 10 | one-dimensional tensor. 11 | 12 | A tuple of :obj:`(unique_tensor, unique_indices)` is returned, where the 13 | :obj:`unique_indices` are the indices of the elements in the original input 14 | tensor. Note that :obj:`unique_indices` is not guaranteed to be stable on 15 | GPU. 16 | 17 | Args: 18 | src (:class:`Tensor`): The input tensor. 19 | 20 | :rtype: (:class:`Tensor`, :class:`LongTensor`) 21 | """ 22 | 23 | src = src.contiguous().view(-1) 24 | 25 | if src.is_cuda: 26 | out, perm = torch_unique.unique_cuda.unique(src) 27 | else: 28 | out, perm = np.unique(src.numpy(), return_index=True) 29 | out, perm = torch.from_numpy(out), torch.from_numpy(perm) 30 | return out, perm 31 | --------------------------------------------------------------------------------