-
Notifications
You must be signed in to change notification settings - Fork 8
/
Copy pathquantizers.py
105 lines (99 loc) · 5.28 KB
/
quantizers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
from pqCoding import *
from opqCoding import *
from aqCoding import *
dataFolder = './' # set folder for your quantization models structure
class Quantizer():
def __init__(self):
self.prefix = 'base_'
def trainCodebooks(self, data, params):
pass
def getQuantizationError(self, data, params):
pass
def encodeDataset(self, data, params):
pass
def searchNearestNeighbors(self, data, params, k=10000):
pass
def getCodebooksFilename(self, data, params):
return dataFolder + data.prefix + self.prefix + params.prefix + 'codebooks.dat'
def getCodesFilename(self, data, params):
return dataFolder + data.prefix + self.prefix + params.prefix + 'code.dat'
class PqQuantizer(Quantizer):
def __init__(self, threadsCount, itCount):
self.prefix = 'pq_'
self.threadsCount = threadsCount
self.itCount = itCount
def trainCodebooks(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
learnCodebooksPQ(data.learnFilename, \
data.dim, params.M, params.K, \
data.learnPointsCount, \
codebooksFilename, self.threadsCount, \
self.itCount)
def getQuantizationError(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
return getQuantizationErrorPQ(data.testFilename, data.dim, data.testPointsCount, codebooksFilename)
def encodeDataset(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
codeFilename = self.getCodesFilename(data, params)
encodeDatasetPQ(data.baseFilename, data.basePointsCount, \
codebooksFilename, codeFilename, self.threadsCount)
def searchNearestNeighbors(self, data, params, k):
codebooksFilename = self.getCodebooksFilename(data, params)
codeFilename = self.getCodesFilename(data, params)
return searchNearestNeighborsPQ(codeFilename, codebooksFilename, \
data.queriesFilename,
data.queriesCount, k, self.threadsCount)
class OpqQuantizer(Quantizer):
def __init__(self, threadsCount, itCount):
self.prefix = 'opq_'
self.threadsCount = threadsCount
self.itCount = itCount
def trainCodebooks(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
learnCodebooksOPQ(data.learnFilename, \
data.learnPointsCount, \
data.dim, params.M, params.K, \
codebooksFilename, self.itCount)
def getQuantizationError(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
return getQuantizationErrorOPQ(codebooksFilename, \
data.testFilename, \
data.testPointsCount)
def encodeDataset(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
codeFilename = self.getCodesFilename(data, params)
encodeDatasetOPQ(data.baseFilename, \
data.basePointsCount, \
codebooksFilename, codeFilename, self.threadsCount)
def searchNearestNeighbors(self, data, params, k=10000):
codebooksFilename = self.getCodebooksFilename(data, params)
codeFilename = self.getCodesFilename(data, params)
return searchNearestNeighborsOPQ(codeFilename, codebooksFilename, \
data.queriesFilename, \
data.queriesCount, k, self.threadsCount)
class AqQuantizer(Quantizer):
def __init__(self, threadsCount, itCount):
self.prefix = 'aq_'
self.threadsCount = threadsCount
self.itCount = itCount
def trainCodebooks(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
learnCodebooksAQ(data.learnFilename, data.dim, \
params.M, params.K, \
data.learnPointsCount, codebooksFilename, \
self.threadsCount, self.itCount)
def getQuantizationError(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
return getQuantizationErrorAQ(data.testFilename, data.dim, \
data.testPointsCount, codebooksFilename, self.threadsCount)
def encodeDataset(self, data, params):
codebooksFilename = self.getCodebooksFilename(data, params)
codeFilename = self.getCodesFilename(data, params)
encodeDatasetAQ(data.baseFilename, data.basePointsCount, \
codebooksFilename, codeFilename, self.threadsCount)
def searchNearestNeighbors(self, data, params, k=10000):
codebooksFilename = self.getCodebooksFilename(data, params)
codeFilename = self.getCodesFilename(data, params)
return searchNearestNeighborsAQ(codeFilename, codebooksFilename, \
data.queriesFilename, data.queriesCount, \
k, self.threadsCount)