faiss/contrib/torch/quantization.py

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""
This contrib module contains Pytorch code for quantization.
"""

import numpy as np
import torch
import faiss

from faiss.contrib import torch_utils


class Quantizer:

    def __init__(self, d, code_size):
        self.d = d
        self.code_size = code_size

    def train(self, x):
        pass

    def encode(self, x):
        pass

    def decode(self, x):
        pass


class VectorQuantizer(Quantizer):

    def __init__(self, d, k):
        code_size = int(torch.ceil(torch.log2(k) / 8))
        Quantizer.__init__(d, code_size)
        self.k = k

    def train(self, x):
        pass


class ProductQuantizer(Quantizer):

    def __init__(self, d, M, nbits):
        code_size = int(torch.ceil(M * nbits / 8))
        Quantizer.__init__(d, code_size)
        self.M = M
        self.nbits = nbits

    def train(self, x):
        pass