{"id":16402745,"url":"https://github.com/demoriarty/torchpq","last_synced_at":"2025-04-05T07:05:33.168Z","repository":{"id":48726422,"uuid":"337198876","full_name":"DeMoriarty/TorchPQ","owner":"DeMoriarty","description":"Approximate nearest neighbor search with product quantization on GPU in pytorch and cuda","archived":false,"fork":false,"pushed_at":"2023-12-12T23:17:26.000Z","size":16237,"stargazers_count":211,"open_issues_count":6,"forks_count":20,"subscribers_count":6,"default_branch":"main","last_synced_at":"2024-10-12T05:47:08.499Z","etag":null,"topics":["cuda","nearest-neighbor-search","pytorch"],"latest_commit_sha":null,"homepage":"","language":"Cuda","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/DeMoriarty.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-02-08T20:20:35.000Z","updated_at":"2024-09-30T18:48:36.000Z","dependencies_parsed_at":"2023-12-13T00:27:35.744Z","dependency_job_id":"f1c34e95-fd35-434e-ab66-d8cd577c543e","html_url":"https://github.com/DeMoriarty/TorchPQ","commit_stats":{"total_commits":286,"total_committers":6,"mean_commits":"47.666666666666664","dds":"0.30069930069930073","last_synced_commit":"e6ccc8b33c5c7da01f5314176e49d330d84e9ba8"},"previous_names":[],"tags_count":5,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeMoriarty%2FTorchPQ","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeMoriarty%2FTorchPQ/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeMoriarty%2FTorchPQ/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/DeMoriarty%2FTorchPQ/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/DeMoriarty","download_url":"https://codeload.github.com/DeMoriarty/TorchPQ/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247299832,"owners_count":20916190,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cuda","nearest-neighbor-search","pytorch"],"created_at":"2024-10-11T05:47:08.539Z","updated_at":"2025-04-05T07:05:33.151Z","avatar_url":"https://github.com/DeMoriarty.png","language":"Cuda","funding_links":[],"categories":[],"sub_categories":[],"readme":"# TorchPQ\nTorchPQ is a python library for **Approximate Nearest Neighbor Search** (ANNS) and **Maximum Inner Product Search** (MIPS) on GPU using **Product Quantization** (PQ) algorithm. TorchPQ is implemented mainly with PyTorch, with some extra CUDA kernels to accelerate clustering, indexing and searching.\n\n## Install\n- make sure you have the latest version of PyTorch installed: https://pytorch.org/get-started/locally/\n- install or upgrade to CUDA toolkit 11.0 or greater\n- install a version of CuPy library that matches your CUDA toolkit version\n```\npip install cupy-cuda110\npip install cupy-cuda111\npip install cupy-cuda112\n...\n```\nfor a full list of cupy-cuda versions, please go to [Installation Guide](https://docs.cupy.dev/en/stable/install.html#installing-cupy)\n- install TorchPQ\n```\npip install torchpq\n```\n\n## Quick Start\n### IVFPQ\n**I**n**V**erted **F**ile **P**roduct **Q**uantization (IVFPQ) is a type of ANN search algorithm that is designed to do fast and efficient vector search in million, or even billion scale vector sets. check the [original paper](https://hal.inria.fr/inria-00514462v2/document) for more details.  \n\n#### Training\n```python\nfrom torchpq.index import IVFPQIndex\nimport torch\n\nn_data = 1000000 # number of data points\nd_vector = 128 # dimentionality / number of features\n\nindex = IVFPQIndex(\n  d_vector=d_vector,\n  n_subvectors=64,\n  n_cells=1024,\n  initial_size=2048,\n  distance=\"euclidean\",\n)\n\ntrainset = torch.randn(d_vector, n_data, device=\"cuda:0\")\nindex.train(trainset)\n```\nThere are some important parameters that need to be explained:\n- **d_vector**: dimentionality of input vectors. there are 2 constraints on `d_vector`: (1) it needs to be divisible by `n_subvectors`; (2) it needs to be a multiple of 4.*\n- **n_subvectors**: number of subquantizers, essentially this is the byte size of each quantized vector, 64 byte per vector in the above example.**\n- **n_cells**: number of coarse quantizer clusters\n- **initial_size**: initial capacity assigned to each voronoi cell of coarse quantizer.\n`n_cells * initial_size` is the number of vectors that can be stored initially. if any cell has reached its capacity, that cell will be automatically expanded.\nIf you need to add vectors frequently, a larger value for `initial_size` is recommended.\n\nRemember that the shape of any tensor that contains data points has to be ```[d_vector, n_data]```.\n\n\\* the second constraint could be removed in the future  \n\\*\\* actual byte size would be (n_subvectors+9) bytes, 8 bytes for ID and 1 byte for is_empty\n#### Adding new vectors\n```python\nbaseset = torch.randn(d_vector, n_data, device=\"cuda:0\")\nids = torch.arange(n_data, device=\"cuda\")\nindex.add(baseset, ids=ids)\n```\nEach ID in `ids` needs to be a unique int64 (`torch.long`) value that identifies a vector in `x`.\nif `ids` is not provided, it will be set to `torch.arange(n_data, device=\"cuda\") + previous_max_id`\n\n#### Removing vectors\n```python\nindex.remove(ids=ids)\n```\n`index.remove(ids=ids)` will virtually remove vectors with specified `ids` from storage.\nIt ignores ids that doesn't exist.\n\n#### Topk search\n```python\nindex.n_probe = 32\nn_query = 10000\nqueryset = torch.randn(d_vector, n_query, device=\"cuda:0\")\ntopk_values, topk_ids = index.search(queryset, k=100)\n```\n- when `distance=\"inner\"`, `topk_values` are **inner product** of queries and topk closest data points.\n- when `distance=\"euclidean\"`, `topk_values` are **negative squared L2 distance** between queries and topk closest data points.\n- when `distance=\"manhattan\"`, `topk_values` are **negative L1 distance** between queries and topk closest data points.\n- when `distance=\"cosine\"`, `topk_values` are **cosine similarity** between queries and topk closest data points.\n\n#### Encode and Decode\nyou can use IVFPQ as a vector codec for lossy compression of vectors\n```python\ncode = index.encode(queryset)   # compression\nreconstruction = index.decode(code) # reconstruction\n```\n\n#### Save and Load\nMost of the TorchPQ modules are inherited from `torch.nn.Module`, this means you can save and load them just like a regular pytorch model.\n```python\n# Save to PATH\ntorch.save(index.state_dict(), PATH)\n# Load from PATH\nindex.load_state_dict(torch.load(PATH))\n```\n### Clustering\n#### K-means\n```python\nfrom torchpq.clustering import KMeans\nimport torch\n\nn_data = 1000000 # number of data points\nd_vector = 128 # dimentionality / number of features\nx = torch.randn(d_vector, n_data, device=\"cuda\")\n\nkmeans = KMeans(n_clusters=4096, distance=\"euclidean\")\nlabels = kmeans.fit(x)\n```\nNotice that the shape of the tensor that contains data points has to be ```[d_vector, n_data]```, this is consistant in TorchPQ.\n\n#### Multiple concurrent K-means\nSometimes, we have multiple independent datasets that need to be clustered,\ninstead of running multiple KMeans sequentianlly,\nwe can perform multiple kmeans concurrently with **MultiKMeans**\n```python\nfrom torchpq.clustering import MultiKMeans\nimport torch\n\nn_data = 1000000\nn_kmeans = 16\nd_vector = 64\nx = torch.randn(n_kmeans, d_vector, n_data, device=\"cuda\")\nkmeans = MultiKMeans(n_clusters=256, distance=\"euclidean\")\nlabels = kmeans.fit(x)\n```\n#### Prediction with K-means\n```\nlabels = kmeans.predict(x)\n```\n\n## Benchmarks\n- [K-Means](/benchmark/turing/kmeans/README.md)\n- [Sift1M](/benchmark/turing/sift1m/README.md)\n- [Gist1M](/benchmark/turing/gist1m/README.md)\n- Sift10M (coming soon)\n- Sift100M (coming soon)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdemoriarty%2Ftorchpq","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdemoriarty%2Ftorchpq","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdemoriarty%2Ftorchpq/lists"}