3. 搜索算法与实践

本节导读：掌握FAISS的搜索方法和参数配置，学习如何在保证精度的前提下实现高效的相似性搜索。

学习目标

• 掌握FAISS搜索API的使用方法
• 理解搜索参数的配置技巧
• 学会使用不同的距离度量方法
• 实现批量搜索和并行处理

核心概念

搜索是FAISS的核心功能，它通过高效的算法在海量向量库中快速找到与查询向量最相似的Top-K个向量。理解搜索原理和掌握参数配置是使用FAISS的关键。

搜索API详解

基本搜索方法

单向量搜索

import faiss
import numpy as np

# 创建索引（假设已经构建好）
d = 128
index = faiss.IndexFlatL2(d)  # 使用Flat索引作为示例

# 生成查询向量
query_vector = np.random.random((1, d)).astype('float32')

# 执行搜索
k = 10  # 返回前10个最近邻
distances, indices = index.search(query_vector, k)

print("距离向量:", distances)
print("索引数组:", indices)
​

批量搜索

# 批量搜索多个查询向量
batch_size = 5
query_vectors = np.random.random((batch_size, d)).astype('float32')

distances, indices = index.search(query_vectors, k)
print("批量搜索结果形状:", distances.shape, indices.shape)
​

搜索结果解读

距离矩阵

• distances[i, j]：第i个查询向量的第j个最近邻的距离
• 距离越小表示相似度越高
• 根据索引类型不同，距离可能有不同含义

索引数组

• indices[i, j]：第i个查询向量的第j个最近邻在原始数据中的索引
• 用于定位对应的原始数据或元信息
• 如果为-1表示该位置没有有效结果

搜索参数配置

核心参数说明

nprobe (仅适用于IVF索引)

# 对于IVF索引，nprobe决定搜索时检查的簇数量
index.nprobe = 20  # 检查20个簇，提升精度但降低速度
​

k (返回数量)

k = 10  # 返回前10个最近邻
# 建议根据应用场景选择合适的k值
# 推荐系统通常k=50-100，检索系统通常k=10-20
​

num_threads

# 设置并行线程数
index.num_threads = 4  # 使用4个线程并行搜索
# 通常设置为CPU核心数，不超过8
​

search_k

# 对于IVF索引，可以设置实际搜索的向量数量
# 用于平衡精度和速度的权衡
index.quantizer.search_k = 1000  # 搜索1000个候选向量
​

参数优化策略

精度与速度的权衡

def optimize_search_parameters(index, validation_queries, validation_ground_truth):
    best_nprobe = 1
    best_recall = 0
    
    # 尝试不同的nprobe值
    for nprobe in [1, 5, 10, 20, 50]:
        index.nprobe = nprobe
        
        # 计算召回率
        distances, indices = index.search(validation_queries, 10)
        recall = calculate_recall(indices, validation_ground_truth)
        
        # 记录最佳参数
        if recall > best_recall:
            best_recall = recall
            best_nprobe = nprobe
            
        print(f"nprobe={nprobe}, recall={recall:.4f}")
    
    print(f"最优参数: nprobe={best_nprobe}, recall={best_recall:.4f}")
    return best_nprobe

def calculate_recall(pred_indices, ground_truth, k=10):
    correct = 0
    for i, gt in enumerate(ground_truth):
        correct += len(set(pred_indices[i]) & set(gt[:k]))
    return correct / (len(ground_truth) * k)
​

多种距离度量方法

常用距离度量

L2距离（欧氏距离）

# L2距离索引
index_l2 = faiss.IndexFlatL2(d)

# 适用场景：
# - 最常用的距离度量
# - 适用于大多数向量表示
# - 计算简单，易于理解
​

内积（Inner Product）

# 内积索引
index_ip = faiss.IndexFlatIP(d)

# 适用场景：
# - 适用于归一化向量
# - 语义相似度计算
# - 在推荐系统中常用

# 注意：使用内积前需要将向量归一化
index_ip = faiss.IndexIDMap2(faiss.IndexFlatIP(d))
faiss.normalize_L2(x)  # 向量归一化
​

余弦相似度

# 余弦相似度（通过内积实现）
index_cosine = faiss.IndexFlatIP(d)

# 向量归一化
normalized_vectors = x / np.linalg.norm(x, axis=1, keepdims=True)

# 搜索结果已经是余弦相似度（值越大越相似）
​

马氏距离

# 马氏距离需要预先计算协方差矩阵
# 需要自定义实现或使用其他距离度量
​

距离度量选择指南

def select_distance_metric(vector_type, application_scenario):
    """根据应用场景选择合适的距离度量"""
    
    if vector_type == "normalized":
        return "inner_product"
    elif vector_type == "raw":
        return "l2"
    elif application_scenario == "semantic_similarity":
        return "cosine"
    elif application_scenario == "feature_matching":
        return "l2"
    else:
        return "l2"  # 默认使用L2距离
​

批量搜索与并行处理

批量搜索优化

class BatchSearch:
    def __init__(self, index, batch_size=1000):
        self.index = index
        self.batch_size = batch_size
        
    def search_batch(self, query_vectors, k=10):
        """分批搜索大批量数据"""
        n_queries = query_vectors.shape[0]
        all_distances = []
        all_indices = []
        
        for i in range(0, n_queries, self.batch_size):
            batch_end = min(i + self.batch_size, n_queries)
            batch_queries = query_vectors[i:batch_end]
            
            distances, indices = self.index.search(batch_queries, k)
            all_distances.append(distances)
            all_indices.append(indices)
            
        return np.vstack(all_distances), np.vstack(all_indices)
​

并行处理实现

import concurrent.futures
import threading

class ParallelSearch:
    def __init__(self, index, max_workers=4):
        self.index = index
        self.max_workers = max_workers
        self.lock = threading.Lock()
        
    def search_parallel(self, query_vectors, k=10):
        """使用多线程并行搜索"""
        n_queries = query_vectors.shape[0]
        batch_size = n_queries // self.max_workers
        
        results = []
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = []
            
            for i in range(0, n_queries, batch_size):
                batch_end = min(i + batch_size, n_queries)
                batch_queries = query_vectors[i:batch_end]
                
                future = executor.submit(self.index.search, batch_queries, k)
                futures.append(future)
                
            # 收集结果
            for future in concurrent.futures.as_completed(futures):
                distances, indices = future.result()
                results.append((distances, indices))
                
        # 合并结果
        all_distances = np.vstack([r[0] for r in results])
        all_indices = np.vstack([r[1] for r in results])
        
        return all_distances, all_indices
​

GPU加速搜索

# 使用GPU加速搜索
res = faiss.StandardGpuResources()  # GPU资源

# 将索引复制到GPU
gpu_index = faiss.index_cpu_to_gpu(res, 0, index)

# GPU搜索
distances, indices = gpu_index.search(query_vectors, k)

# 搜索完成后将结果复制回CPU
if faiss.get_num_gpus() > 0:
    distances = distances.copy_to_host()
    indices = indices.copy_to_host()
​

搜索结果解析与后处理

结果验证

def validate_search_results(index, query_vectors, expected_results, k=10):
    """验证搜索结果的正确性"""
    distances, indices = index.search(query_vectors, k)
    
    validation_results = []
    
    for i, query in enumerate(query_vectors):
        # 获取查询向量的最近邻
        nearest_neighbors = indices[i]
        nearest_distances = distances[i]
        
        # 验证结果格式
        assert len(nearest_neighbors) == k
        assert len(nearest_distances) == k
        
        # 验证索引有效性
        valid_indices = [idx for idx in nearest_neighbors if idx != -1]
        assert len(valid_indices) <= k
        
        validation_results.append({
            'query_index': i,
            'nearest_neighbors': nearest_neighbors,
            'distances': nearest_distances,
            'valid_count': len(valid_indices)
        })
    
    return validation_results
​

结果排序和过滤

def post_process_search_results(distances, indices, min_distance=0.0, max_distance=float('inf')):
    """后处理搜索结果"""
    processed_results = []
    
    for i in range(indices.shape[0]):
        # 过滤距离超出范围的条目
        valid_mask = (distances[i] >= min_distance) & (distances[i] <= max_distance)
        filtered_distances = distances[i][valid_mask]
        filtered_indices = indices[i][valid_mask]
        
        # 按距离排序
        sorted_indices = np.argsort(filtered_distances)
        sorted_distances = filtered_distances[sorted_indices]
        sorted_indices = filtered_indices[sorted_indices]
        
        processed_results.append({
            'indices': sorted_indices,
            'distances': sorted_distances
        })
    
    return processed_results
​

缓存机制

import pickle
import os

class SearchCache:
    def __init__(self, cache_dir="search_cache"):
        self.cache_dir = cache_dir
        os.makedirs(cache_dir, exist_ok=True)
        
    def get_cache_key(self, query_vectors, k, search_params):
        """生成缓存键"""
        import hashlib
        query_str = str(query_vectors.tobytes())
        params_str = str(search_params)
        combined = query_str + params_str
        return hashlib.md5(combined.encode()).hexdigest()
        
    def save_to_cache(self, cache_key, distances, indices):
        """保存到缓存"""
        cache_file = os.path.join(self.cache_dir, f"{cache_key}.pkl")
        with open(cache_file, 'wb') as f:
            pickle.dump((distances, indices), f)
            
    def load_from_cache(self, cache_key):
        """从缓存加载"""
        cache_file = os.path.join(self.cache_dir, f"{cache_key}.pkl")
        if os.path.exists(cache_file):
            with open(cache_file, 'rb') as f:
                return pickle.load(f)
        return None
​

完整搜索示例

import faiss
import numpy as np
import time

class FAISSSearchEngine:
    def __init__(self, dimension=128, index_type="IVF"):
        self.dimension = dimension
        self.index_type = index_type
        self.index = None
        self.is_trained = False
        
    def build_index(self, data, nlist=100):
        """构建索引"""
        print(f"构建{self.index_type}索引...")
        
        if self.index_type == "Flat":
            self.index = faiss.IndexFlatL2(self.dimension)
        elif self.index_type == "IVF":
            quantizer = faiss.IndexFlatL2(self.dimension)
            self.index = faiss.IndexIVFFlat(quantizer, self.dimension, nlist)
            
        # 训练索引（需要时）
        if self.index_type == "IVF":
            self.index.train(data)
            self.is_trained = True
            
        # 添加数据到索引
        self.index.add(data)
        print(f"索引构建完成，包含{self.index.ntotal}个向量")
        
    def search(self, queries, k=10, nprobe=10):
        """执行搜索"""
        if self.index is None:
            raise ValueError("索引未构建")
            
        # 设置搜索参数
        if hasattr(self.index, 'nprobe'):
            self.index.nprobe = nprobe
            
        # 执行搜索
        start_time = time.time()
        distances, indices = self.index.search(queries, k)
        search_time = time.time() - start_time
        
        return distances, indices, search_time
        
    def batch_search(self, queries, k=10, batch_size=1000):
        """批量搜索"""
        n_queries = queries.shape[0]
        all_distances = []
        all_indices = []
        total_time = 0
        
        for i in range(0, n_queries, batch_size):
            batch_end = min(i + batch_size, n_queries)
            batch_queries = queries[i:batch_end]
            
            distances, indices, time_used = self.search(batch_queries, k)
            all_distances.append(distances)
            all_indices.append(indices)
            total_time += time_used
            
        return np.vstack(all_distances), np.vstack(all_indices), total_time

# 使用示例
if __name__ == "__main__":
    # 生成示例数据
    dimension = 128
    n_data = 100000
    n_queries = 1000
    
    data = np.random.random((n_data, dimension)).astype('float32')
    queries = np.random.random((n_queries, dimension)).astype('float32')
    
    # 创建搜索引擎
    engine = FAISSSearchEngine(dimension=128, index_type="IVF")
    
    # 构建索引
    engine.build_index(data, nlist=100)
    
    # 执行搜索
    distances, indices, search_time = engine.search(queries, k=10)
    
    print(f"搜索完成，耗时: {search_time:.2f}秒")
    print(f"平均每个查询耗时: {search_time/n_queries*1000:.2f}毫秒")
​

本章小结

本章详细介绍了FAISS的搜索算法和实践方法，包括搜索API的使用、参数配置、多种距离度量方法、批量搜索和并行处理，以及搜索结果的后处理技术。通过完整的示例代码，读者可以掌握如何在实际项目中使用FAISS进行高效的相似性搜索。下一章我们将探讨FAISS的高级特性和优化技术。

关键词：搜索算法, 参数配置, 距离度量, 批量搜索, 并行处理
难度：进阶
预计阅读：45分钟