mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Tia Lim
0f0f728166
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5349 # this is a no-op change Ran `black -l 80` over all tracked Python files under `fbcode/faiss/`, using line length 80 to match faiss's documented coding style (`CONTRIBUTING.md`: "80 character line length (both for C++ and Python)"). 172 files were reformatted; 28 were already compliant. This is a pure-formatting change. `black` runs its AST-equivalence safety check on every file, so no identifiers, string/numeric values, operators, or control flow were altered — only whitespace, indentation, line wrapping, blank lines, trailing commas, and string-quote normalization (single to double quotes). `fbcode/faiss/` is intentionally excluded from Meta's standard Python autoformatter (`arc f` / `arc lint -a` report no formatter for these paths) because it mirrors to public GitHub, so the Python had drifted out of any consistent style. The C++ is already clang-format-clean (0 of 807 files change), so this effort is Python-only. This diff brings the Python to canonical `black` formatting. Excluded `demos/index_pq_flat_separate_codes_from_codebook.py`: it contains a Jupyter `!rm` shell-magic line and is not valid Python, so `black` cannot parse it. This is the base of a 3-diff stack: (1) this black reformat; (2) `.flake8` alignment (max-line-length 80 + black-compatible ignores E203/W503); (3) genuine lint fixes (F401/F403/F841/B011). Reviewed By: mnorris11 Differential Revision: D109736099 fbshipit-source-id: 9a054635766ed11aa3760a788395babd3d198abf
98 lines
2.7 KiB
Python
98 lines
2.7 KiB
Python
# Copyright (c) Meta Platforms, Inc. and 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 torch
|
|
import faiss
|
|
import math
|
|
from faiss.contrib.torch import clustering
|
|
|
|
# the kmeans can produce both torch and numpy centroids
|
|
|
|
|
|
class Quantizer:
|
|
|
|
def __init__(self, d, code_size):
|
|
"""
|
|
d: dimension of vectors
|
|
code_size: nb of bytes of the code (per vector)
|
|
"""
|
|
self.d = d
|
|
self.code_size = code_size
|
|
|
|
def train(self, x):
|
|
"""
|
|
takes a n-by-d array and performs training
|
|
"""
|
|
pass
|
|
|
|
def encode(self, x):
|
|
"""
|
|
takes a n-by-d float array, encodes to an n-by-code_size uint8 array
|
|
"""
|
|
pass
|
|
|
|
def decode(self, codes):
|
|
"""
|
|
takes a n-by-code_size uint8 array, returns a n-by-d array
|
|
"""
|
|
pass
|
|
|
|
|
|
class VectorQuantizer(Quantizer):
|
|
|
|
def __init__(self, d, k):
|
|
|
|
code_size = int(math.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):
|
|
"""M: number of subvectors, d%M == 0
|
|
nbits: number of bits that each vector is encoded into
|
|
"""
|
|
assert d % M == 0
|
|
assert nbits == 8 # todo: implement other nbits values
|
|
code_size = int(math.ceil(M * nbits / 8))
|
|
Quantizer.__init__(self, d, code_size)
|
|
self.M = M
|
|
self.nbits = nbits
|
|
self.code_size = code_size
|
|
|
|
def train(self, x):
|
|
nc = 2**self.nbits
|
|
sd = self.d // self.M
|
|
dev = x.device
|
|
dtype = x.dtype
|
|
self.codebook = torch.zeros((self.M, nc, sd), device=dev, dtype=dtype)
|
|
for m in range(self.M):
|
|
xsub = x[:, m * self.d // self.M : (m + 1) * self.d // self.M]
|
|
data = clustering.DatasetAssign(xsub.contiguous())
|
|
self.codebook[m] = clustering.kmeans(2**self.nbits, data)
|
|
|
|
def encode(self, x):
|
|
codes = torch.zeros((x.shape[0], self.code_size), dtype=torch.uint8)
|
|
for m in range(self.M):
|
|
xsub = x[:, m * self.d // self.M : (m + 1) * self.d // self.M]
|
|
_, I = faiss.knn(xsub.contiguous(), self.codebook[m], 1)
|
|
codes[:, m] = I.ravel()
|
|
return codes
|
|
|
|
def decode(self, codes):
|
|
idxs = [codes[:, m].long() for m in range(self.M)]
|
|
vectors = [self.codebook[m, idxs[m], :] for m in range(self.M)]
|
|
stacked_vectors = torch.stack(vectors, dim=1)
|
|
cbd = self.codebook.shape[-1]
|
|
x_rec = stacked_vectors.reshape(-1, cbd * self.M)
|
|
return x_rec
|