Skip to content
Permalink
Browse files

fix(encoder): fix bug in vlad encoder

  • Loading branch information...
Larryjianfeng committed Sep 9, 2019
1 parent 1ba4e11 commit ddf13ff1ecf014599a32031e2071f0b57336fe55
Showing with 8 additions and 5 deletions.
  1. +1 −1 gnes/encoder/numeric/vlad.py
  2. +7 −4 tests/test_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
@@ -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))

0 comments on commit ddf13ff

Please sign in to comment.
You can’t perform that action at this time.