mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
clean up serialization tests (#4700)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4700 A pass over the serialization / deserialization tests, so that they don't use temp files when not needed. Reviewed By: pankajsingh88, junjieqi Differential Revision: D87615610 fbshipit-source-id: 3e6366b938d41e083781745179ff19ecf5247ba2
This commit is contained in:
committed by
meta-codesync[bot]
parent
ac2e3abe38
commit
1ea99d8073
+24
-79
@@ -10,8 +10,6 @@ from __future__ import absolute_import, division, print_function, unicode_litera
|
||||
import numpy as np
|
||||
import unittest
|
||||
import faiss
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
|
||||
def make_binary_dataset(d, nb, nt, nq):
|
||||
@@ -37,21 +35,12 @@ class TestBinaryFlat(unittest.TestCase):
|
||||
index = faiss.IndexBinaryFlat(d)
|
||||
index.add(self.xb)
|
||||
D, I = index.search(self.xq, 3)
|
||||
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
|
||||
|
||||
fd, tmpnam = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index_binary(index, tmpnam)
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
index2 = faiss.read_index_binary(tmpnam)
|
||||
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
finally:
|
||||
os.remove(tmpnam)
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
|
||||
class TestBinaryIVF(unittest.TestCase):
|
||||
@@ -76,20 +65,12 @@ class TestBinaryIVF(unittest.TestCase):
|
||||
index.add(self.xb)
|
||||
D, I = index.search(self.xq, 3)
|
||||
|
||||
fd, tmpnam = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index_binary(index, tmpnam)
|
||||
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
|
||||
|
||||
index2 = faiss.read_index_binary(tmpnam)
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
finally:
|
||||
os.remove(tmpnam)
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
|
||||
class TestObjectOwnership(unittest.TestCase):
|
||||
@@ -109,16 +90,10 @@ class TestObjectOwnership(unittest.TestCase):
|
||||
index = faiss.IndexBinaryFlat(d)
|
||||
index.add(self.xb)
|
||||
|
||||
fd, tmpnam = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index_binary(index, tmpnam)
|
||||
# this is the output of read_index_binary (==> checks ownership)
|
||||
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
|
||||
|
||||
index2 = faiss.read_index_binary(tmpnam)
|
||||
|
||||
assert index2.thisown
|
||||
finally:
|
||||
os.remove(tmpnam)
|
||||
assert index2.thisown
|
||||
|
||||
|
||||
class TestBinaryFromFloat(unittest.TestCase):
|
||||
@@ -140,21 +115,11 @@ class TestBinaryFromFloat(unittest.TestCase):
|
||||
index.add(self.xb)
|
||||
D, I = index.search(self.xq, 3)
|
||||
|
||||
fd, tmpnam = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index_binary(index, tmpnam)
|
||||
|
||||
index2 = faiss.read_index_binary(tmpnam)
|
||||
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
finally:
|
||||
os.remove(tmpnam)
|
||||
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
class TestBinaryHNSW(unittest.TestCase):
|
||||
|
||||
@@ -174,20 +139,12 @@ class TestBinaryHNSW(unittest.TestCase):
|
||||
index.add(self.xb)
|
||||
D, I = index.search(self.xq, 3)
|
||||
|
||||
fd, tmpnam = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index_binary(index, tmpnam)
|
||||
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
|
||||
|
||||
index2 = faiss.read_index_binary(tmpnam)
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
finally:
|
||||
os.remove(tmpnam)
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
def test_ivf_hnsw(self):
|
||||
d = self.xq.shape[1] * 8
|
||||
@@ -200,21 +157,9 @@ class TestBinaryHNSW(unittest.TestCase):
|
||||
index.add(self.xb)
|
||||
D, I = index.search(self.xq, 3)
|
||||
|
||||
fd, tmpnam = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index_binary(index, tmpnam)
|
||||
index2 = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
|
||||
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
index2 = faiss.read_index_binary(tmpnam)
|
||||
|
||||
D2, I2 = index2.search(self.xq, 3)
|
||||
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
finally:
|
||||
os.remove(tmpnam)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
assert (I2 == I).all()
|
||||
assert (D2 == D).all()
|
||||
|
||||
+6
-22
@@ -6,8 +6,6 @@
|
||||
|
||||
import unittest
|
||||
import time
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import faiss
|
||||
@@ -587,16 +585,9 @@ class TestAQFastScan(unittest.TestCase):
|
||||
index.add(ds.get_database())
|
||||
D1, I1 = index.search(ds.get_queries(), 1)
|
||||
|
||||
fd, fname = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index(index, fname)
|
||||
index2 = faiss.read_index(fname)
|
||||
D2, I2 = index2.search(ds.get_queries(), 1)
|
||||
np.testing.assert_array_equal(I1, I2)
|
||||
finally:
|
||||
if os.path.exists(fname):
|
||||
os.unlink(fname)
|
||||
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
||||
D2, I2 = index2.search(ds.get_queries(), 1)
|
||||
np.testing.assert_array_equal(I1, I2)
|
||||
|
||||
def test_io(self):
|
||||
self.subtest_io('LSQ4x4fs_Nlsq2x4')
|
||||
@@ -685,16 +676,9 @@ class TestPAQFastScan(unittest.TestCase):
|
||||
index.add(ds.get_database())
|
||||
D1, I1 = index.search(ds.get_queries(), 1)
|
||||
|
||||
fd, fname = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index(index, fname)
|
||||
index2 = faiss.read_index(fname)
|
||||
D2, I2 = index2.search(ds.get_queries(), 1)
|
||||
np.testing.assert_array_equal(I1, I2)
|
||||
finally:
|
||||
if os.path.exists(fname):
|
||||
os.unlink(fname)
|
||||
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
||||
D2, I2 = index2.search(ds.get_queries(), 1)
|
||||
np.testing.assert_array_equal(I1, I2)
|
||||
|
||||
def test_io(self):
|
||||
self.subtest_io('PLSQ2x3x4fs_Nlsq2x4')
|
||||
|
||||
@@ -4,9 +4,7 @@
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
|
||||
import os
|
||||
import unittest
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import faiss
|
||||
@@ -821,18 +819,11 @@ class TestIVFAQFastScan(unittest.TestCase):
|
||||
index = faiss.index_factory(d, factory_str)
|
||||
index.train(ds.get_train())
|
||||
index.add(ds.get_database())
|
||||
D1, I1 = index.search(ds.get_queries(), 1)
|
||||
_, I1 = index.search(ds.get_queries(), 1)
|
||||
|
||||
fd, fname = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index(index, fname)
|
||||
index2 = faiss.read_index(fname)
|
||||
D2, I2 = index2.search(ds.get_queries(), 1)
|
||||
np.testing.assert_array_equal(I1, I2)
|
||||
finally:
|
||||
if os.path.exists(fname):
|
||||
os.unlink(fname)
|
||||
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
||||
_, I2 = index2.search(ds.get_queries(), 1)
|
||||
np.testing.assert_array_equal(I1, I2)
|
||||
|
||||
def test_io(self):
|
||||
self.subtest_io('IVF16,LSQ4x4fs_Nlsq2x4')
|
||||
@@ -929,16 +920,9 @@ class TestIVFPAQFastScan(unittest.TestCase):
|
||||
index.add(ds.get_database())
|
||||
D1, I1 = index.search(ds.get_queries(), 1)
|
||||
|
||||
fd, fname = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index(index, fname)
|
||||
index2 = faiss.read_index(fname)
|
||||
D2, I2 = index2.search(ds.get_queries(), 1)
|
||||
np.testing.assert_array_equal(I1, I2)
|
||||
finally:
|
||||
if os.path.exists(fname):
|
||||
os.unlink(fname)
|
||||
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
||||
D2, I2 = index2.search(ds.get_queries(), 1)
|
||||
np.testing.assert_array_equal(I1, I2)
|
||||
|
||||
def test_io(self):
|
||||
self.subtest_io('IVF16,PLSQ2x3x4fsr_Nlsq2x4')
|
||||
|
||||
@@ -15,7 +15,6 @@ Paper: https://www.arxiv.org/pdf/2510.00566
|
||||
"""
|
||||
|
||||
import unittest
|
||||
import tempfile
|
||||
import os
|
||||
|
||||
import faiss
|
||||
|
||||
@@ -8,8 +8,6 @@
|
||||
import numpy as np
|
||||
import unittest
|
||||
import faiss
|
||||
import tempfile
|
||||
import os
|
||||
|
||||
from common_faiss_tests import get_dataset_2
|
||||
|
||||
@@ -280,14 +278,7 @@ class TestNSG(unittest.TestCase):
|
||||
return knn_graph
|
||||
|
||||
def subtest_io_and_clone(self, index, Dnsg, Insg):
|
||||
fd, tmpfile = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index(index, tmpfile)
|
||||
index2 = faiss.read_index(tmpfile)
|
||||
finally:
|
||||
if os.path.exists(tmpfile):
|
||||
os.unlink(tmpfile)
|
||||
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
||||
|
||||
Dnsg2, Insg2 = index2.search(self.xq, 1)
|
||||
np.testing.assert_array_equal(Dnsg2, Dnsg)
|
||||
@@ -306,8 +297,6 @@ class TestNSG(unittest.TestCase):
|
||||
|
||||
def subtest_add(self, build_type, thresh, metric=faiss.METRIC_L2):
|
||||
d = self.xq.shape[1]
|
||||
metrics = {faiss.METRIC_L2: 'L2',
|
||||
faiss.METRIC_INNER_PRODUCT: 'IP'}
|
||||
|
||||
flat_index = faiss.IndexFlat(d, metric)
|
||||
flat_index.add(self.xb)
|
||||
@@ -381,8 +370,6 @@ class TestNSG(unittest.TestCase):
|
||||
def test_reset(self):
|
||||
"""test IndexNSG.reset()"""
|
||||
d = self.xq.shape[1]
|
||||
metrics = {faiss.METRIC_L2: 'L2',
|
||||
faiss.METRIC_INNER_PRODUCT: 'IP'}
|
||||
|
||||
metric = faiss.METRIC_L2
|
||||
flat_index = faiss.IndexFlat(d, metric)
|
||||
@@ -546,14 +533,7 @@ class TestNNDescent(unittest.TestCase):
|
||||
self.assertGreaterEqual(recalls, 450) # 462
|
||||
|
||||
# do some IO tests
|
||||
fd, tmpfile = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index(index, tmpfile)
|
||||
index2 = faiss.read_index(tmpfile)
|
||||
finally:
|
||||
if os.path.exists(tmpfile):
|
||||
os.unlink(tmpfile)
|
||||
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
||||
|
||||
D2, I2 = index2.search(self.xq, 1)
|
||||
np.testing.assert_array_equal(D2, D)
|
||||
@@ -592,10 +572,6 @@ class TestNNDescentKNNG(unittest.TestCase):
|
||||
self.subtest(32, 10, faiss.METRIC_INNER_PRODUCT)
|
||||
|
||||
def subtest(self, d, K, metric):
|
||||
metric_names = {faiss.METRIC_L1: 'L1',
|
||||
faiss.METRIC_L2: 'L2',
|
||||
faiss.METRIC_INNER_PRODUCT: 'IP'}
|
||||
|
||||
nb = 1000
|
||||
_, xb, _ = get_dataset_2(d, 0, nb, 0)
|
||||
|
||||
|
||||
@@ -4,7 +4,6 @@
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
""" more elaborate that test_index.py """
|
||||
from __future__ import absolute_import, division, print_function
|
||||
|
||||
import numpy as np
|
||||
import unittest
|
||||
@@ -193,13 +192,7 @@ class TestRemove(unittest.TestCase):
|
||||
assert False, 'should have raised an exception'
|
||||
|
||||
# while we are there, let's test I/O as well...
|
||||
fd, tmpnam = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index_binary(index, tmpnam)
|
||||
index = faiss.read_index_binary(tmpnam)
|
||||
finally:
|
||||
os.remove(tmpnam)
|
||||
index = faiss.deserialize_index_binary(faiss.serialize_index_binary(index))
|
||||
|
||||
assert index.reconstruct(1004)[0] == 104
|
||||
try:
|
||||
@@ -469,14 +462,7 @@ class TestIVFFlatDedup(unittest.TestCase):
|
||||
check_ref_knn_with_draws(Dref, Iref, Dnew, Inew)
|
||||
|
||||
# test I/O
|
||||
fd, tmpfile = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
try:
|
||||
faiss.write_index(index_new, tmpfile)
|
||||
index_st = faiss.read_index(tmpfile)
|
||||
finally:
|
||||
if os.path.exists(tmpfile):
|
||||
os.unlink(tmpfile)
|
||||
index_st = faiss.deserialize_index(faiss.serialize_index(index_new))
|
||||
Dst, Ist = index_st.search(xq, 20)
|
||||
|
||||
check_ref_knn_with_draws(Dnew, Inew, Dst, Ist)
|
||||
|
||||
+9
-11
@@ -359,21 +359,19 @@ class Test_IO_PQ(unittest.TestCase):
|
||||
|
||||
try:
|
||||
faiss.write_ProductQuantizer(index.pq, fname)
|
||||
|
||||
read_pq = faiss.read_ProductQuantizer(fname)
|
||||
|
||||
self.assertEqual(index.pq.M, read_pq.M)
|
||||
self.assertEqual(index.pq.nbits, read_pq.nbits)
|
||||
self.assertEqual(index.pq.dsub, read_pq.dsub)
|
||||
self.assertEqual(index.pq.ksub, read_pq.ksub)
|
||||
np.testing.assert_array_equal(
|
||||
faiss.vector_to_array(index.pq.centroids),
|
||||
faiss.vector_to_array(read_pq.centroids)
|
||||
)
|
||||
|
||||
finally:
|
||||
if os.path.exists(fname):
|
||||
os.unlink(fname)
|
||||
self.assertEqual(index.pq.M, read_pq.M)
|
||||
self.assertEqual(index.pq.nbits, read_pq.nbits)
|
||||
self.assertEqual(index.pq.dsub, read_pq.dsub)
|
||||
self.assertEqual(index.pq.ksub, read_pq.ksub)
|
||||
np.testing.assert_array_equal(
|
||||
faiss.vector_to_array(index.pq.centroids),
|
||||
faiss.vector_to_array(read_pq.centroids)
|
||||
)
|
||||
|
||||
|
||||
|
||||
class Test_IO_IndexLSH(unittest.TestCase):
|
||||
|
||||
@@ -15,7 +15,6 @@ Paper: https://www.arxiv.org/pdf/2510.00566
|
||||
"""
|
||||
|
||||
import unittest
|
||||
import tempfile
|
||||
import os
|
||||
|
||||
import faiss
|
||||
|
||||
Reference in New Issue
Block a user