昇腾社区首页
中文
注册

AscendIndexIVFSQ with PCAR

大库算法AscendIndexIVFSQ可以根据一组数据结合量化与IVF的方法进行训练,生成合适的量化和分桶函数,对于输入的float32的特征向量,AscendIndexIVFSQ对其量化为Int8类型的特征向量并将其分配至指定的分桶,并在执行向量比对的时候,选取距离待比对向量最近的nprobe个分桶,将其中Int8类型的向量反量化为原始的特征向量执行后续的计算。PCAR为mxIndex默认提供的对应的降维算法,通过主成分分析的方法能够将高维向量进行进一步的降维,从而换取更高的检索速度和存储容量,同时mxIndex也支持用户自定义的神经网络的降维方式。典型带PCAR降维的AscendIndexIVFSQ的样例参考如下。

#include <faiss/ascend/AscendIndexIVFSQ.h>
#include <faiss/ascend/AscendIndexPreTransform.h>
#include <iostream>
#include <random>

using namespace std;

int main(int argc, char **argv)
{
    const size_t dim = 128;
    const size_t dimReduced = 64;
    const size_t ntotal = 10000;
    const size_t ncentroids = 1024;
    vector<float> data(dim * ntotal);
    for (int i = 0; i < data.size(); i++) {
        data[i] = drand48();
    }

    const size_t k = 100;
    const size_t searchNum = 100;
    const size_t nprobe = 24;

    vector<float> dist(k * searchNum);
    vector<long> indices(k * searchNum);

    faiss::ascend::AscendIndexPreTransform *index = nullptr;
    try {
        faiss::ascend::AscendIndexIVFSQConfig conf({0}, 768);
        faiss::ascend::AscendIndexIVFSQ *subIndex = new faiss::ascend::AscendIndexIVFSQ(
            dimReduced, ncentroids, faiss::ScalarQuantizer::QuantizerType::QT_8bit, faiss::METRIC_INNER_PRODUCT, false, conf);
        subIndex->verbose = true;
        subIndex->setNumProbes(nprobe);

        index = new faiss::ascend::AscendIndexPreTransform(subIndex);
        index->verbose = true;
        index->prependTransform<faiss::ascend::AscendPCAMatrix>(dim, dimReduced, 0.0f, true);
        index->train(ntotal, data.data());
        index->add(ntotal, data.data());
        index->search(searchNum, data.data(), k, dist.data(), indices.data());
    } catch (...) {
        cout << "Exception caught!" << endl;
        delete index;
        return -1;
    }
    cout << "Search finished successfully" << endl;
    delete index;
    return 0;
}