fix(encoder): fix bug in vlad encoder

gnes-ai · Sep 9, 2019 · ddf13ff · ddf13ff
1 parent 1ba4e11
commit ddf13ff
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/gnes/encoder/numeric/vlad.py b/gnes/encoder/numeric/vlad.py
@@ -26,7 +26,7 @@ class VladEncoder(BaseNumericEncoder):
     batch_size = 2048
 
     def __init__(self, num_clusters: int,
-                 using_faiss_pred: True,
+                 using_faiss_pred: bool=True,
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.num_clusters = num_clusters

diff --git a/tests/test_vlad.py b/tests/test_vlad.py
@@ -18,9 +18,12 @@ def test_vlad_train(self):
         model = VladEncoder(20)
         model.train(self.mock_train_data)
         self.assertEqual(model.centroids.shape, (20, 128))
-        model.dump(self.dump_path)
-
-    def test_vlad_encode(self):
-        model = VladEncoder.load(self.dump_path)
         v = model.encode(self.mock_eval_data)
         self.assertEqual(v.shape, (2, 2560))
+
+    def test_vlad_dump_load(self):
+        model = VladEncoder(20)
+        model.train(self.mock_train_data)
+        model.dump(self.dump_path)
+        model_new = VladEncoder.load(self.dump_path)
+        self.assertEqual(model_new.centroids.shape, (20, 128))