From 8dfd51d3a76a04cda2bcbbd224cb89b2393351d3 Mon Sep 17 00:00:00 2001 From: Matthijs Douze Date: Fri, 28 Aug 2020 17:49:25 -0700 Subject: [PATCH] Moved pytorch interop code to contrib Summary: The pytorch interop code was in a test until now. However, it is better if people can rely on it to be updated when the API is updated. Therefore, we move it into contrib. Also added a README.md Reviewed By: wickedfoo Differential Revision: D23392962 fbshipit-source-id: 9b7c0e388a7ea3c0b73dc0018322138f49191673 --- contrib/README.md | 45 ++++++++++++ contrib/pytorch_tensors.py | 102 +++++++++++++++++++++++++++ faiss/gpu/test/test_pytorch_faiss.py | 97 +------------------------ 3 files changed, 148 insertions(+), 96 deletions(-) create mode 100644 contrib/README.md create mode 100644 contrib/pytorch_tensors.py diff --git a/contrib/README.md b/contrib/README.md new file mode 100644 index 000000000..1c77e1c37 --- /dev/null +++ b/contrib/README.md @@ -0,0 +1,45 @@ + +# The contrib modules + +The contrib directory contains helper modules for Faiss for various tasks. + +## Code structure + +The contrib directory gets compiled in the module faiss.contrib. +Note that although some of the modules may depend on additional modules (eg. GPU Faiss, pytorch, hdf5), they are not necessarily compiled in to avoid adding dependencies. It is the user's responsibility to provide them. + +In contrib, we are progressively dropping python2 support. + +## List of contrib modules + +### rpc.py + +A very simple Remote Procedure Call library, where function parameters and results are pickled, for use with client_server.py + +### client_server.py + +The server handles requests to a Faiss index. The client calls the remote index. +This is mainly to shard datasets over several machines, see [Distributd index](https://github.com/facebookresearch/faiss/wiki/Indexes-that-do-not-fit-in-RAM#distributed-index) + +### ondisk.py + +Encloses the main logic to merge indexes into an on-disk index. +See [On-disk storage](https://github.com/facebookresearch/faiss/wiki/Indexes-that-do-not-fit-in-RAM#on-disk-storage) + +### exhaustive_search.py + +Computes the ground-truth search results for a dataset that possibly does not fit in RAM. Uses GPU if available. +Tested in `tests/test_contrib.TestComputeGT` + +### gpu.py + +(requires GPU Faiss) + +Interoperability functions for pytorch and Faiss: pass GPU data without copying back to CPU. +Tested in `gpu/test/test_pytorch_faiss` + +### datasets.py + +(may require h5py) + +Defintion of how to access data for some standard datsets. diff --git a/contrib/pytorch_tensors.py b/contrib/pytorch_tensors.py new file mode 100644 index 000000000..ae559582f --- /dev/null +++ b/contrib/pytorch_tensors.py @@ -0,0 +1,102 @@ +import faiss +import torch + +def swig_ptr_from_FloatTensor(x): + """ gets a Faiss SWIG pointer from a pytorch trensor (on CPU or GPU) """ + assert x.is_contiguous() + assert x.dtype == torch.float32 + return faiss.cast_integer_to_float_ptr( + x.storage().data_ptr() + x.storage_offset() * 4) + +def swig_ptr_from_LongTensor(x): + """ gets a Faiss SWIG pointer from a pytorch trensor (on CPU or GPU) """ + assert x.is_contiguous() + assert x.dtype == torch.int64, 'dtype=%s' % x.dtype + return faiss.cast_integer_to_long_ptr( + x.storage().data_ptr() + x.storage_offset() * 8) + + + +def search_index_pytorch(index, x, k, D=None, I=None): + """call the search function of an index with pytorch tensor I/O (CPU + and GPU supported)""" + assert x.is_contiguous() + n, d = x.size() + assert d == index.d + + if D is None: + D = torch.empty((n, k), dtype=torch.float32, device=x.device) + else: + assert D.size() == (n, k) + + if I is None: + I = torch.empty((n, k), dtype=torch.int64, device=x.device) + else: + assert I.size() == (n, k) + torch.cuda.synchronize() + xptr = swig_ptr_from_FloatTensor(x) + Iptr = swig_ptr_from_LongTensor(I) + Dptr = swig_ptr_from_FloatTensor(D) + index.search_c(n, xptr, + k, Dptr, Iptr) + torch.cuda.synchronize() + return D, I + + +def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None, + metric=faiss.METRIC_L2): + """search xq in xb, without building an index""" + assert xb.device == xq.device + + nq, d = xq.size() + if xq.is_contiguous(): + xq_row_major = True + elif xq.t().is_contiguous(): + xq = xq.t() # I initially wrote xq:t(), Lua is still haunting me :-) + xq_row_major = False + else: + raise TypeError('matrix should be row or column-major') + + xq_ptr = swig_ptr_from_FloatTensor(xq) + + nb, d2 = xb.size() + assert d2 == d + if xb.is_contiguous(): + xb_row_major = True + elif xb.t().is_contiguous(): + xb = xb.t() + xb_row_major = False + else: + raise TypeError('matrix should be row or column-major') + xb_ptr = swig_ptr_from_FloatTensor(xb) + + if D is None: + D = torch.empty(nq, k, device=xb.device, dtype=torch.float32) + else: + assert D.shape == (nq, k) + assert D.device == xb.device + + if I is None: + I = torch.empty(nq, k, device=xb.device, dtype=torch.int64) + else: + assert I.shape == (nq, k) + assert I.device == xb.device + + D_ptr = swig_ptr_from_FloatTensor(D) + I_ptr = swig_ptr_from_LongTensor(I) + + args = faiss.GpuDistanceParams() + args.metric = metric + args.k = k + args.dims = d + args.vectors = xb_ptr + args.vectorsRowMajor = xb_row_major + args.numVectors = nb + args.queries = xq_ptr + args.queriesRowMajor = xq_row_major + args.numQueries = nq + args.outDistances = D_ptr + args.outIndices = I_ptr + faiss.bfKnn(res, args) + + return D, I diff --git a/faiss/gpu/test/test_pytorch_faiss.py b/faiss/gpu/test/test_pytorch_faiss.py index ba5582b35..ccc459870 100644 --- a/faiss/gpu/test/test_pytorch_faiss.py +++ b/faiss/gpu/test/test_pytorch_faiss.py @@ -10,102 +10,7 @@ import unittest import faiss import torch -def swig_ptr_from_FloatTensor(x): - assert x.is_contiguous() - assert x.dtype == torch.float32 - return faiss.cast_integer_to_float_ptr( - x.storage().data_ptr() + x.storage_offset() * 4) - -def swig_ptr_from_LongTensor(x): - assert x.is_contiguous() - assert x.dtype == torch.int64, 'dtype=%s' % x.dtype - return faiss.cast_integer_to_long_ptr( - x.storage().data_ptr() + x.storage_offset() * 8) - - - -def search_index_pytorch(index, x, k, D=None, I=None): - """call the search function of an index with pytorch tensor I/O (CPU - and GPU supported)""" - assert x.is_contiguous() - n, d = x.size() - assert d == index.d - - if D is None: - D = torch.empty((n, k), dtype=torch.float32, device=x.device) - else: - assert D.size() == (n, k) - - if I is None: - I = torch.empty((n, k), dtype=torch.int64, device=x.device) - else: - assert I.size() == (n, k) - torch.cuda.synchronize() - xptr = swig_ptr_from_FloatTensor(x) - Iptr = swig_ptr_from_LongTensor(I) - Dptr = swig_ptr_from_FloatTensor(D) - index.search_c(n, xptr, - k, Dptr, Iptr) - torch.cuda.synchronize() - return D, I - - -def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None, - metric=faiss.METRIC_L2): - assert xb.device == xq.device - - nq, d = xq.size() - if xq.is_contiguous(): - xq_row_major = True - elif xq.t().is_contiguous(): - xq = xq.t() # I initially wrote xq:t(), Lua is still haunting me :-) - xq_row_major = False - else: - raise TypeError('matrix should be row or column-major') - - xq_ptr = swig_ptr_from_FloatTensor(xq) - - nb, d2 = xb.size() - assert d2 == d - if xb.is_contiguous(): - xb_row_major = True - elif xb.t().is_contiguous(): - xb = xb.t() - xb_row_major = False - else: - raise TypeError('matrix should be row or column-major') - xb_ptr = swig_ptr_from_FloatTensor(xb) - - if D is None: - D = torch.empty(nq, k, device=xb.device, dtype=torch.float32) - else: - assert D.shape == (nq, k) - assert D.device == xb.device - - if I is None: - I = torch.empty(nq, k, device=xb.device, dtype=torch.int64) - else: - assert I.shape == (nq, k) - assert I.device == xb.device - - D_ptr = swig_ptr_from_FloatTensor(D) - I_ptr = swig_ptr_from_LongTensor(I) - - args = faiss.GpuDistanceParams() - args.metric = metric - args.k = k - args.dims = d - args.vectors = xb_ptr - args.vectorsRowMajor = xb_row_major - args.numVectors = nb - args.queries = xq_ptr - args.queriesRowMajor = xq_row_major - args.numQueries = nq - args.outDistances = D_ptr - args.outIndices = I_ptr - faiss.bfKnn(res, args) - - return D, I +from faiss.contrib.pytorch_tensors import search_index_pytorch, search_raw_array_pytorch def to_column_major(x): if hasattr(torch, 'contiguous_format'):