mirror of
https://github.com/facebookresearch/faiss.git
synced 2025-06-03 21:54:02 +08:00
dd72e4121d
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/3608 This is a straightforward implementation of QINCo in CPU Faiss, with encoding and decoding capabilities (not training). For this, we translate a simplified version of some torch classes: - tensors, restricted to 2D and int32 + float32 - Linear and Embedding layer Then the QINCoStep and QINCo can just be defined as C++ objects that are copy-constructable. There is some plumbing required in the wrapping layers to support the integration. Pytroch tensors are converted to numpy for getting / setting them in C++. Reviewed By: asadoughi Differential Revision: D59132952 fbshipit-source-id: eea4856507a5b7c5f219efcf8d19fe56944df088
57 lines
1.8 KiB
C++
57 lines
1.8 KiB
C++
/**
|
|
* 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.
|
|
*/
|
|
|
|
#include <faiss/IndexNeuralNetCodec.h>
|
|
#include <faiss/impl/FaissAssert.h>
|
|
#include <faiss/utils/hamming.h>
|
|
|
|
namespace faiss {
|
|
|
|
/*********************************************************
|
|
* IndexNeuralNetCodec implementation
|
|
*********************************************************/
|
|
|
|
IndexNeuralNetCodec::IndexNeuralNetCodec(
|
|
int d,
|
|
int M,
|
|
int nbits,
|
|
MetricType metric)
|
|
: IndexFlatCodes((M * nbits + 7) / 8, d, metric), M(M), nbits(nbits) {
|
|
is_trained = false;
|
|
}
|
|
|
|
void IndexNeuralNetCodec::train(idx_t n, const float* x) {
|
|
FAISS_THROW_MSG("Training not implemented in C++, use Pytorch");
|
|
}
|
|
|
|
void IndexNeuralNetCodec::sa_encode(idx_t n, const float* x, uint8_t* codes)
|
|
const {
|
|
nn::Tensor2D x_tensor(n, d, x);
|
|
nn::Int32Tensor2D codes_tensor = net->encode(x_tensor);
|
|
pack_bitstrings(n, M, nbits, codes_tensor.data(), codes, code_size);
|
|
}
|
|
|
|
void IndexNeuralNetCodec::sa_decode(idx_t n, const uint8_t* codes, float* x)
|
|
const {
|
|
nn::Int32Tensor2D codes_tensor(n, M);
|
|
unpack_bitstrings(n, M, nbits, codes, code_size, codes_tensor.data());
|
|
nn::Tensor2D x_tensor = net->decode(codes_tensor);
|
|
memcpy(x, x_tensor.data(), d * n * sizeof(float));
|
|
}
|
|
|
|
/*********************************************************
|
|
* IndexQINeuralNetCodec implementation
|
|
*********************************************************/
|
|
|
|
IndexQINCo::IndexQINCo(int d, int M, int nbits, int L, int h, MetricType metric)
|
|
: IndexNeuralNetCodec(d, M, nbits, metric),
|
|
qinco(d, 1 << nbits, L, M, h) {
|
|
net = &qinco;
|
|
}
|
|
|
|
} // namespace faiss
|