8d776e6453
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/1484 This diff allows for native usage of PyTorch tensors for Faiss indexes on both CPU and GPU. It is currently only implemented in this diff for things that inherit from `faiss.Index`, which covers the non-binary indices, and it patches the same functions on `faiss.Index` that were also covered by `__init__.py` for numpy interoperability. There must be uniformity among the inputs: if any array input is a Torch tensor, then all array inputs must be Torch tensors. Similarly, if any array input is a numpy ndarray, then all array inputs must be numpy ndarrays. If `faiss.contrib.torch_utils` is imported, it ensures that `import faiss` has already been performed to patch all of the functions using the base `__init__.py` numpy wrappers, and then patches the following functions again: ``` add add_with_ids assign train search remove_ids reconstruct reconstruct_n range_search update_vectors search_and_reconstruct sa_encode sa_decode ``` to allow usage of PyTorch CPU tensors, and additionally PyTorch GPU tensors if the index being used is on the GPU. numpy functionality is still available when `faiss.contrib.torch_utils` is imported; we pass through to the original patched numpy function when we detect numpy inputs. In addition, to allow for better (asynchronous) GPU usage without requiring the CPU to be involved, all of these functions which construct tensors/arrays for output now take optional arguments for storage (numpy or torch.Tensor) to be provided that will contain the output data. `range_search` is the only exception to this, as the size of the output data is indeterminate. The eventual GPU implementation will likely require the user to provide a maximum cap on the output size, and allow that to be passed instead. If the optional pre-allocated output values are presented by the user, they are used; otherwise, new return ndarray / Tensors are constructed as before and used for the return. If this feature were not provided on the GPU, then every execution would be completely serial as we would depend upon the CPU to allocate GPU memory before every operation. Instead, now this can function much like NN graph execution on the GPU, assuming that all of the data requirements are pre-allocated, so the execution will run at the full speed of the GPU and not be stalled sequentially launching kernels. This diff also exposes the `GpuResources` shared_ptr object owned by a GPU index. This is required for pytorch GPU so that we can perform proper stream ordering in Faiss with respect to the current pytorch stream. So, Faiss indices now perform more or less as any NN operation in Torch does. Note, however, that a Faiss index has its own setting on current device, and if the pytorch GPU tensor inputs are resident on a different device than what the Faiss index expects, a cross-device copy will be initiated. I may choose to make this an error in the future and require matching device to device. This diff also found a bug when passing GPU data directly to `train()` for `GpuIndexIVFFlat` and `GpuIndexIVFScalarQuantizer`, as I guess we never tested passing GPU data directly to these functions before. `GpuIndexIVFPQ` was doing the right thing however. The assign function is now also implemented on the GPU as well, and is now marked `const` to be in line with the `search` function. Also added better checking of non-contiguous inputs for both Torch tensors and numpy ndarrays. Updated the `knn_gpu` function with a base implementation always present that allows for usage of numpy arrays, which is overridden when `torch_utils` is imported to allow torch usage. This supports row/column major layout, float32/float16 data and int64/int32 indices for both numpy and torch. Reviewed By: mdouze Differential Revision: D24299400 fbshipit-source-id: b4f117b9c120bd1ad83e7702087051ab7b303b29
Faiss
Faiss is a library for efficient similarity search and clustering of dense vectors. It contains algorithms that search in sets of vectors of any size, up to ones that possibly do not fit in RAM. It also contains supporting code for evaluation and parameter tuning. Faiss is written in C++ with complete wrappers for Python/numpy. Some of the most useful algorithms are implemented on the GPU. It is developed by Facebook AI Research.
NEWS
NEW: version 1.6.3 (2020-03-27) IndexBinaryHash, GPU support for alternative distances.
NEW: version 1.6.1 (2019-11-29) bugfix.
NEW: version 1.6.0 (2019-10-15) code structure reorg, support for codec interface.
NEW: version 1.5.3 (2019-06-24) fix performance regression in IndexIVF.
NEW: version 1.5.2 (2019-05-27) the license was relaxed to MIT from BSD+Patents. Read LICENSE for details.
NEW: version 1.5.0 (2018-12-19) GPU binary flat index and binary HNSW index
NEW: version 1.4.0 (2018-08-30) no more crashes in pure Python code
NEW: version 1.3.0 (2018-07-12) support for binary indexes
NEW: latest commit (2018-02-22) supports on-disk storage of inverted indexes, see demos/demo_ondisk_ivf.py
NEW: latest commit (2018-01-09) includes an implementation of the HNSW indexing method, see benchs/bench_hnsw.py
NEW: there is now a Facebook public discussion group for Faiss users at https://www.facebook.com/groups/faissusers/
NEW: on 2017-07-30, the license on Faiss was relaxed to BSD from CC-BY-NC. Read LICENSE for details.
Introduction
Faiss contains several methods for similarity search. It assumes that the instances are represented as vectors and are identified by an integer, and that the vectors can be compared with L2 (Euclidean) distances or dot products. Vectors that are similar to a query vector are those that have the lowest L2 distance or the highest dot product with the query vector. It also supports cosine similarity, since this is a dot product on normalized vectors.
Most of the methods, like those based on binary vectors and compact quantization codes, solely use a compressed representation of the vectors and do not require to keep the original vectors. This generally comes at the cost of a less precise search but these methods can scale to billions of vectors in main memory on a single server.
The GPU implementation can accept input from either CPU or GPU memory. On a server with GPUs, the GPU indexes can be used a drop-in replacement for the CPU indexes (e.g., replace IndexFlatL2 with GpuIndexFlatL2) and copies to/from GPU memory are handled automatically. Results will be faster however if both input and output remain resident on the GPU. Both single and multi-GPU usage is supported.
Building
The library is mostly implemented in C++, with optional GPU support provided via CUDA, and an optional Python interface. The CPU version requires a BLAS library. It compiles with a Makefile and can be packaged in a docker image. See INSTALL.md for details.
How Faiss works
Faiss is built around an index type that stores a set of vectors, and provides a function to search in them with L2 and/or dot product vector comparison. Some index types are simple baselines, such as exact search. Most of the available indexing structures correspond to various trade-offs with respect to
- search time
- search quality
- memory used per index vector
- training time
- need for external data for unsupervised training
The optional GPU implementation provides what is likely (as of March 2017) the fastest exact and approximate (compressed-domain) nearest neighbor search implementation for high-dimensional vectors, fastest Lloyd's k-means, and fastest small k-selection algorithm known. The implementation is detailed here.
Full documentation of Faiss
The following are entry points for documentation:
- the full documentation, including a tutorial, a FAQ and a troubleshooting section can be found on the wiki page
- the doxygen documentation gives per-class information
- to reproduce results from our research papers, Polysemous codes and Billion-scale similarity search with GPUs, refer to the benchmarks README. For Link and code: Fast indexing with graphs and compact regression codes, see the link_and_code README
Authors
The main authors of Faiss are:
- Hervé Jégou initiated the Faiss project and wrote its first implementation
- Matthijs Douze implemented most of the CPU Faiss
- Jeff Johnson implemented all of the GPU Faiss
- Lucas Hosseini implemented the binary indexes
Reference
Reference to cite when you use Faiss in a research paper:
@article{JDH17,
title={Billion-scale similarity search with GPUs},
author={Johnson, Jeff and Douze, Matthijs and J{\'e}gou, Herv{\'e}},
journal={arXiv preprint arXiv:1702.08734},
year={2017}
}
Join the Faiss community
For public discussion of Faiss or for questions, there is a Facebook group at https://www.facebook.com/groups/faissusers/
We monitor the issues page of the repository. You can report bugs, ask questions, etc.
License
Faiss is MIT-licensed.