Faiss(12)：python接口faiss.py文件分析

最新推荐文章于 2024-07-05 03:46:41 发布

翔底

最新推荐文章于 2024-07-05 03:46:41 发布

阅读量2k

点赞数 1

分类专栏： Faiss

本文链接：https://blog.csdn.net/rangfei/article/details/108949735

版权

Faiss 专栏收录该内容

17 篇文章 42 订阅

订阅专栏

1. 前言

本篇笔记主要分析faiss code下的python接口文件——faiss.py的工作流程以及内容。

2. faiss.py分析

2.1 导入文件

在faiss code 编译完成后，在python目录下执行make/make install命令后，会在该文件下产生faiss.py, faiss/, faiss.egg-info/等文件和目录，其中faiss.py和faiss/目录下的__init__.py内容一样，但是在应用程序中import faiss时，导入的是__init__.py文件。

2.2 执行流程

faiss.py在导入时，会执行以下工作：

导入依赖模块

import numpy as np
import sys
import inspect
import pdb
import platform
import subprocess
import logging

创建一个log文件,这里__name__是"faiss"

logger = logging.getLogger(__name__)

根据指令集加载swigfaiss包

try:
    instr_set = instruction_set()
    if instr_set == "AVX2":
        logger.info("Loading faiss with AVX2 support.")
        from .swigfaiss_avx2 import *
    else:
        logger.info("Loading faiss.")
        from .swigfaiss import *

except ImportError:
    # we import * so that the symbol X can be accessed as faiss.X
    logger.info("Loading faiss.")
    from .swigfaiss import *

在instruction_set方法中会调用platform.system()和subprocess.check_output来获取当前平台的系统信息和CPU的指令集信息。我使用的平台返回的是Linux和"AVX2"。

记录版本信息

__version__ = "%d.%d.%d" % (FAISS_VERSION_MAJOR,
                            FAISS_VERSION_MINOR,
                            FAISS_VERSION_PATCH)

上述信息定义在Index.h文件中，我使用的版本是1.6.1

替换clustering的train方法

handle_Clustering()

def handle_Clustering():
    def replacement_train(self, x, index):
        assert x.flags.contiguous
        n, d = x.shape
        assert d == self.d
        self.train_c(n, swig_ptr(x), index)
    replace_method(Clustering, 'train', replacement_train)


def replace_method(the_class, name, replacement, ignore_missing=False):
    try:
        orig_method = getattr(the_class, name)
    except AttributeError:
        if ignore_missing:
            return
        raise
    if orig_method.__name__ == 'replacement_' + name:
        # replacement was done in parent class
        return
    setattr(the_class, name + '_c', orig_method)
    setattr(the_class, name, replacement)

将clustering class中的train方法名称替换为replacement的名称。
replace_method方法中：

the_class：类名，文件中直接使用的类定义在swigfaiss.py中，即这里能调用的类都必须是在该文件中有定义的。
name：函数名，即要重新替换的index的方法
replacement：替换后的方法，这里传入方法的地址，在faiss.py中针对index的部分函数如train, add等定义了一套python下的方法，主要是添加了"replacement_"前缀。所以，在应用程序中调用上述方法时会运行到替换后的方法中来。如index.add --> replacement_add

替换 ProductQuantizer 和 ScalarQuantizer 的方法

handle_Quantizer(ProductQuantizer)
handle_Quantizer(ScalarQuantizer)

def handle_Quantizer(the_class):

    def replacement_train(self, x):
        n, d = x.shape
        assert d == self.d
        self.train_c(n, swig_ptr(x))

    def replacement_compute_codes(self, x):
        n, d = x.shape
        assert d == self.d
        codes = np.empty((n, self.code_size), dtype='uint8')
        self.compute_codes_c(swig_ptr(x), swig_ptr(codes), n)
        return codes

    def replacement_decode(self, codes):
        n, cs = codes.shape
        assert cs == self.code_size
        x = np.empty((n, self.d), dtype='float32')
        self.decode_c(swig_ptr(codes), swig_ptr(x), n)
        return x

    replace_method(the_class, 'train', replacement_train)
    replace_method(the_class, 'compute_codes', replacement_compute_codes)
    replace_method(the_class, 'decode', replacement_decode)

这两个类是量化器相关的，这里会替换三个函数： train, compute_codes, decode。

替换MatrixStats的初始化方法

def handle_MatrixStats(the_class):
    original_init = the_class.__init__

    def replacement_init(self, m):
        assert len(m.shape) == 2
        original_init(self, m.shape[0], m.shape[1], swig_ptr(m))

    the_class.__init__ = replacement_init

handle_MatrixStats(MatrixStats)

MatrixStats用于报告数据集相关的一些统计信息并对其进行注释。

替换已经导入的index的部分方法

this_module = sys.modules[__name__]

for symbol in dir(this_module):
    obj = getattr(this_module, symbol)
    # print symbol, isinstance(obj, (type, types.ClassType))
    if inspect.isclass(obj):
        the_class = obj
        if issubclass(the_class, Index):
            handle_Index(the_class)

        if issubclass(the_class, IndexBinary):
            handle_IndexBinary(the_class)

        if issubclass(the_class, VectorTransform):
            handle_VectorTransform(the_class)

        if issubclass(the_class, AutoTuneCriterion):
            handle_AutoTuneCriterion(the_class)

        if issubclass(the_class, ParameterSpace):
            handle_ParameterSpace(the_class)

首先sys.modules返回所有已经导入的faiss相关的module，然后根据symbol名称判断是否为class，如果是，则根据该class的父类类型依次替换部分方法，具体的替换过程在handle_xxx中实现。

对类或方法添加参数引用

add_ref_in_constructor(IndexIVFFlat, 0)
add_ref_in_constructor(IndexIVFFlatDedup, 0)
add_ref_in_constructor(IndexPreTransform, {2: [0, 1], 1: [0]})
add_ref_in_method(IndexPreTransform, 'prepend_transform', 0)
add_ref_in_constructor(IndexIVFPQ, 0)
add_ref_in_constructor(IndexIVFPQR, 0)
add_ref_in_constructor(Index2Layer, 0)
add_ref_in_constructor(Level1Quantizer, 0)
add_ref_in_constructor(IndexIVFScalarQuantizer, 0)
add_ref_in_constructor(IndexIDMap, 0)
add_ref_in_constructor(IndexIDMap2, 0)
add_ref_in_constructor(IndexHNSW, 0)
add_ref_in_method(IndexShards, 'add_shard', 0)
add_ref_in_method(IndexBinaryShards, 'add_shard', 0)
add_ref_in_constructor(IndexRefineFlat, 0)
add_ref_in_constructor(IndexBinaryIVF, 0)
add_ref_in_constructor(IndexBinaryFromFloat, 0)
add_ref_in_constructor(IndexBinaryIDMap, 0)
add_ref_in_constructor(IndexBinaryIDMap2, 0)

add_ref_in_method(IndexReplicas, 'addIndex', 0)
add_ref_in_method(IndexBinaryReplicas, 'addIndex', 0)

# seems really marginal...
# remove_ref_from_method(IndexReplicas, 'removeIndex', 0)

if hasattr(this_module, 'GpuIndexFlat'):
    # handle all the GPUResources refs
    add_ref_in_function('index_cpu_to_gpu', 0)
    add_ref_in_constructor(GpuIndexFlat, 0)
    add_ref_in_constructor(GpuIndexFlatIP, 0)
    add_ref_in_constructor(GpuIndexFlatL2, 0)
    add_ref_in_constructor(GpuIndexIVFFlat, 0)
    add_ref_in_constructor(GpuIndexIVFScalarQuantizer, 0)
    add_ref_in_constructor(GpuIndexIVFPQ, 0)
    add_ref_in_constructor(GpuIndexBinaryFlat, 0)

add_ref_in_constructor()：对arg1的类添加arg2的参数引用
add_ref_in_method()：对arg1类中arg2的方法添加arg3的参数引用
add_ref_in_function()：对arg1的函数添加arg2的参数引用，其中arg1的函数是独立于类的函数。

替换MapLong2Long的方法

replace_method(MapLong2Long, 'add', replacement_map_add)
replace_method(MapLong2Long, 'search_multiple', replacement_map_search_multiple)

2.3 定义的方法接口

faiss.py除了在import 时会自动执行上述操作，在文件中还定义了一些独立的函数供用户应用程序调用。

1. GPU相关操作

index_cpu_to_gpu_multiple_py()
此方法的作用是将index从CPU中拷贝到所有的GPU中，为GPU索引和资源构建c++向量。

方法调用了swigfaiss.py的index_cpu_to_gpu_multiple()，该函数拷贝操作的具体实现，初始定义在gpu/GpuCloner.cpp中。

def index_cpu_to_gpu_multiple_py(resources, index, co=None):
    """builds the C++ vectors for the GPU indices and the
    resources. Handles the common case where the resources are assigned to
    the first len(resources) GPUs"""
    vres = GpuResourcesVector()
    vdev = IntVector()
    for i, res in enumerate(resources):
        vdev.push_back(i)
        vres.push_back(res)
    index = index_cpu_to_gpu_multiple(vres, vdev, index, co)
    index.referenced_objects = resources
    return index

resources：资源列表
index： CPU中索引实例
return index： GPU中的索引实例

index_cpu_to_all_gpus()
此方法的作用与同样是拷贝index，但是接口更简单，更易于应用程序调用，其内部通过调用index_cpu_to_gpu_multiple_py实现。

def index_cpu_to_all_gpus(index, co=None, ngpu=-1):
    if ngpu == -1:
        ngpu = get_num_gpus()
    res = [StandardGpuResources() for i in range(ngpu)]
    index2 = index_cpu_to_gpu_multiple_py(res, index, co)
    return index2

index: 要拷贝的CPU 索引实例
co：拷贝的参数，默认为0
ngpu： gpu个数，默认为-1，此时由系统判断
index2：拷贝后GPU中的索引实例

2. numpy数组和std::vector转换操作

vector_to_array()
将c++的std::vector容器转换成np的数组

def vector_to_array(v):
    """ convert a C++ vector to a numpy array """
    classname = v.__class__.__name__
    assert classname.endswith('Vector')
    dtype = np.dtype(vector_name_map[classname[:-6]])
    a = np.empty(v.size(), dtype=dtype)
    if v.size() > 0:
        memcpy(swig_ptr(a), v.data(), a.nbytes)
    return a

# 其中vector_name_map表示了std::vector和python中数据类型的映射：
vector_name_map = {
    'Float': 'float32',
    'Byte': 'uint8',
    'Char': 'int8',
    'Uint64': 'uint64',
    'Long': 'int64',
    'Int': 'int32',
    'Double': 'float64'
    }

v：要转换的c++类型的数组
a：转换后的numpy风格的列表

vector_float_to_array()
同vector_to_array()

def vector_float_to_array(v):
    return vector_to_array(v)

copy_array_to_vector()
将numpy的列表拷贝到c++的容器中

def copy_array_to_vector(a, v):
    """ copy a numpy array to a vector """
    n, = a.shape
    classname = v.__class__.__name__
    assert classname.endswith('Vector')
    dtype = np.dtype(vector_name_map[classname[:-6]])
    assert dtype == a.dtype, (
        'cannot copy a %s array to a %s (should be %s)' % (
            a.dtype, classname, dtype))
    v.resize(n)
    if n > 0:
        memcpy(v.data(), swig_ptr(a), a.nbytes)

3. 数据操作相关的接口

kmin(array, k)
返回数组array中每一行的k个最小值

array：要搜索的二维数组
k：要查找的最小值的个数

返回值：

D ：返回<m, k>的二维数组，类型为int
I ：返回<m, k>的二维数组，类型为float32

kmax(array, k)
返回数组array中每一行的k个最大值，其他同kmin()

pairwise_distances(xq, xb, mt=METRIC_L2, metric_arg=0)
计算两组向量之间的整个成对距离矩阵

xq: 第一个二维向量
xb: 第二个二维向量
mt: 计算的距离类型，默认为L2
metric_arg：距离类型参数

返回值，dis: 一个<m1, m2>的二维向量，其中m1为xq的行数，m2为xb的行数

rand/randint/randn

rand：产生浮点型的随机数
randint：产生int64型随机数
randn：从正态分布中产生folat32型随机数

eval_intersection(I1, I2)
两个结果表的每行之间相交的大小

normalize_L2(x)

def normalize_L2(x):
    fvec_renorm_L2(x.shape[1], x.shape[0], swig_ptr(x))

4. kmeans类

class Kmeans:
    """shallow wrapper around the Clustering object. The important method is train()."""

    def __init__(self, d, k, **kwargs):
        """d: input dimension, k: nb of centroids. Additional
         parameters are passed on the ClusteringParameters object,
         including niter=25, verbose=False, spherical = False
        """
        self.d = d
        self.k = k
        self.gpu = False
        self.cp = ClusteringParameters()
        for k, v in kwargs.items():
            if k == 'gpu':
                self.gpu = v
            else:
                # if this raises an exception, it means that it is a non-existent field
                getattr(self.cp, k)
                setattr(self.cp, k, v)
        self.centroids = None

    def train(self, x):
        n, d = x.shape
        assert d == self.d
        clus = Clustering(d, self.k, self.cp)
        if self.cp.spherical:
            self.index = IndexFlatIP(d)
        else:
            self.index = IndexFlatL2(d)
        if self.gpu:
            if self.gpu == True:
                ngpu = -1
            else:
                ngpu = self.gpu
            self.index = index_cpu_to_all_gpus(self.index, ngpu=ngpu)
        clus.train(x, self.index)
        centroids = vector_float_to_array(clus.centroids)
        self.centroids = centroids.reshape(self.k, d)
        self.obj = vector_float_to_array(clus.obj)
        return self.obj[-1] if self.obj.size > 0 else 0.0

    def assign(self, x):
        assert self.centroids is not None, "should train before assigning"
        self.index.reset()
        self.index.add(self.centroids)
        D, I = self.index.search(x, 1)
        return D.ravel(), I.ravel()

5. 索引与数组的转换

serialize_index()
将索引转换成uint8型数组

def serialize_index(index):
    """ convert an index to a numpy uint8 array  """
    writer = VectorIOWriter()
    write_index(index, writer)
    return vector_to_array(writer.data)

deserialize_index()
使用数组生成索引

def deserialize_index(data):
    reader = VectorIOReader()
    copy_array_to_vector(data, reader.data)
    return read_index(reader)

3. 总结

导入faiss.py后可以直接生成各类Index的实例，其中参数与c++中初始定义的相同，但是并不是所有的函数都可以直接调用，取决于该文件中是否替换该类型的index。
faiss.py本身与c++ core不关联，它通过swigfaiss而形成完整的index的操作；
swigfaiss.py是由SWIG将C++编写的软件嵌入联接成的python语言的接口模块。要继续分析修改c++代码（添加函数和功能）如何反应在python应用程序中，还需要研究SWIG软件的编译规则，并修改之。