Faiss: Can PyTorch Tensors be used as the index?

[moskomule]

Hi, I'm trying to use FAISS with PyTorch (both are the latest versions available via conda install -c pytorch ...). In https://github.com/facebookresearch/faiss/blob/master/gpu/test/test_pytorch_faiss.py, the index is constructed as follows.

xq = faiss.randn(nq * d, 1234).reshape(nq, d)
xb = faiss.randn(nb * d, 1235).reshape(nb, d)
res = faiss.StandardGpuResources()
index = faiss.GpuIndexFlatIP(res, d)
index.add(xb)

Is it possible to use PyTorch's Tensor as xb above? Thank you.

---

At least index.add(torch_tensor) doesn't work.

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-10-12010aa60bdb> in <module>()
      1 index = faiss.IndexFlatL2(d)
----> 2 index.add(xb_torch)

~/.miniconda/lib/python3.6/site-packages/faiss/__init__.py in replacement_add(self, x)
     95 
     96     def replacement_add(self, x):
---> 97         assert x.flags.contiguous
     98         n, d = x.shape
     99         assert d == self.d

AttributeError: 'Tensor' object has no attribute 'flags'

[beauby]

No, you'll need to convert your tensor to a numpy array.

[moskomule]

Thank you. 😢

[beauby]

@moskomule Note that, as mentioned here, it is just a matter of doing index.add(xb.numpy()) instead of index.add(xb).

[moskomule]

I see. But in my case, xbis a gpu tensor so I'm just afraid the bottleneck (GPU->CPU->GPU) and extra memory consumption.

[mdouze]

See here on how to pass pytorch GPU tensors without copying them:
https://github.com/facebookresearch/faiss/blob/master/gpu/test/test_pytorch_faiss.py

[moskomule]

@mdouze Thank you. You mean tensor.storage().data_ptr()? Still unclear to me how to use it as the index.

[gemfield]

The link above is incorrect now. Access https://github.com/facebookresearch/faiss/blob/master/faiss/gpu/test/test_pytorch_faiss.py instead.
You can also have a look at https://github.com/DeepVAC/deepvac/blob/master/deepvac/syszux_feature_vector.py, the class NamesPathsClsFeatureVectorByFaissPytorch may help.

[skei0]

https://github.com/facebookresearch/faiss/blob/master/faiss/gpu/test/test_pytorch_faiss.py and https://github.com/facebookresearch/faiss/blob/master/gpu/test/test_pytorch_faiss.py go to 404 error for me.
Is it still possible to use Pytorch GPU tensor?

[wickedfoo]

@skei0 yes, import faiss.contrib.torch_utils (https://github.com/facebookresearch/faiss/blob/master/contrib/torch_utils.py) and you can simply pass a pytorch tensor as xb.