faiss使用实例，快速入门

吴大炮

已于 2024-05-21 18:48:56 修改

阅读量703

点赞数 5

分类专栏：学习记录文章标签： python

于 2024-05-21 18:31:54 首次发布

本文链接：https://blog.csdn.net/weixin_44645198/article/details/139099431

版权

学习记录专栏收录该内容

30 篇文章 14 订阅

订阅专栏

faiss使用实例，快速入门

基本步骤

导入FAISS库
创建一个FAISS索引。
将数据添加到索引中。
运行最近邻搜索

实例

import numpy as np
import faiss

# 生成一些随机数据
dimension = 128  # 数据的维度
db_size = 10000  # 数据库的大小
np.random.seed(1234)  # 为了结果的可重复性
db_vectors = np.random.random((db_size, dimension)).astype('float32')

# 创建一个FAISS索引
index = faiss.IndexFlatL2(dimension)

# 将数据添加到索引中
index.add(db_vectors)

# 生成一些查询向量
n_query = 10  # 查询的数量
query_vectors = np.random.random((n_query, dimension)).astype('float32')

# 运行最近邻搜索
k = 4  # 我们想要搜索的最近邻的数量
distances, indices = index.search(query_vectors, k)

# 打印结果
print(indices)  # 每行包含一个查询的最近邻的索引
print(distances)  # 每行包含一个查询的最近邻的距离

在这个示例中，我们首先生成了一些随机数据，然后创建了一个FAISS索引，并将数据添加到索引中。然后，我们生成了一些查询向量，并运行了最近邻搜索。最后，我们打印了搜索结果，包括每个查询的最近邻的索引和距离。

如何在FAISS中使用GPU进行最近邻搜索？

这段代码的主要目标是使用FAISS库在GPU上创建、训练、保存、加载和查询一个向量索引。以下是每个步骤的详细说明：

生成一些随机数据：这段代码首先生成了一些随机的浮点数向量，这些向量将被用于训练和添加到FAISS索引中。
创建一个FAISS索引：然后，代码创建了一个FAISS索引。这个索引是IndexIVFFlat类型，这是一种需要训练的索引。
将索引转移到GPU上：为了加速计算，代码将索引转移到了GPU上。
训练索引：然后，代码使用生成的随机数据训练了索引。
将数据添加到索引中：训练后，代码将随机数据添加到了索引中。
保存索引到文件：然后，代码将索引从GPU转移到CPU，并将其保存到文件中，以便以后使用。
从文件加载索引：代码从文件中加载了索引，并将其转移到GPU上。
使用加载的索引进行搜索：最后，代码生成了一些新的随机查询向量，并使用加载的索引对这些查询进行了搜索。搜索的结果是每个查询的最近邻的索引和距离。

总的来说，这段代码展示了如何使用FAISS库在GPU上创建、训练、保存、加载和查询一个向量索引。

import numpy as np
import faiss

# 生成一些随机数据
dimension = 128  # 数据的维度
db_size = 10000  # 数据库的大小
np.random.seed(1234)  # 为了结果的可重复性
db_vectors = np.random.random((db_size, dimension)).astype('float32')

# 创建一个FAISS索引
nlist = 100  # 聚类中心的数量
kmeans_index = faiss.IndexFlatL2(dimension)
index = faiss.IndexIVFFlat(kmeans_index, dimension, nlist)

# 将索引转移到GPU上
res = faiss.StandardGpuResources()  # 创建一个GPU资源对象
gpu_index = faiss.index_cpu_to_gpu(res, 0, index)  # 将索引转移到GPU上

# 训练索引
gpu_index.train(db_vectors)

# 将数据添加到索引中
gpu_index.add(db_vectors)

# 保存索引到文件
# 将索引从GPU转移到CPU
cpu_index = faiss.index_gpu_to_cpu(gpu_index)

# 将索引写入文件
faiss.write_index(cpu_index, 'trained_index_file.index')

# 在下次测试的时候，你可以从文件加载索引
# 从文件中读取索引
cpu_index = faiss.read_index('trained_index_file.index')

# 将索引转移到GPU
gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)

# 然后你可以使用加载的索引进行搜索
n_query = 10  # 查询的数量
query_vectors = np.random.random((n_query, dimension)).astype('float32')
k = 4  # 我们想要搜索的最近邻的数量
distances, indices = gpu_index.search(query_vectors, k)

# 打印结果
print(indices)  # 每行包含一个查询的最近邻的索引
print(distances)  # 每行包含一个查询的最近邻的距离

在FAISS中，"训练"通常是指为某些类型的索引（例如，量化索引）学习一组聚类中心。这些聚类中心用于将高维向量量化为更紧凑的形式，以节省存储空间并加速搜索。

对于IndexFlatL2这样的简单索引，训练过程实际上是不必要的，因为它们不使用量化。但是，对于IndexIVFFlat、IndexIVFPQ等需要量化的索引，训练过程是必要的。

在你的代码片段中，_train方法调用了FAISS索引的train方法。这意味着如果索引需要训练（例如，它是一个IndexIVFFlat或IndexIVFPQ索引），那么这个方法将执行训练过程。如果索引不需要训练（例如，它是一个IndexFlatL2索引），那么这个方法将不会有任何效果。

总的来说，训练的目的是为了使FAISS索引能够更有效地存储和搜索向量。

faiss实现以图搜图

import os
import numpy as np
from PIL import Image
import torch
from torchvision import models, transforms
import faiss

# 图像预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 加载预训练的ResNet模型
model = models.resnet50(pretrained=True)
model = model.eval()

# 提取图像特征
def extract_features(img_path):
    img = Image.open(img_path)
    img = transform(img)
    img = img.unsqueeze(0)
    with torch.no_grad():
        features = model(img)
    return features.numpy()

# 创建一个FAISS索引
dimension = 1000  # ResNet模型的输出维度
index = faiss.IndexFlatL2(dimension)

# 从图像目录中提取特征并添加到索引中
image_dir = 'path_to_your_image_directory'
for img_name in os.listdir(image_dir):
    img_path = os.path.join(image_dir, img_name)
    features = extract_features(img_path)
    index.add(features)

# 保存索引到文件
faiss.write_index(index, 'image_index.index')

# 从文件加载索引
index = faiss.read_index('image_index.index')

# 查询最相似的图像
query_img_path = 'path_to_your_query_image'
query_features = extract_features(query_img_path)
k = 10  # 我们想要搜索的最近邻的数量
distances, indices = index.search(query_features, k)

# 打印结果
print(indices)  # 每行包含一个查询的最近邻的索引
print(distances)  # 每行包含一个查询的最近邻的距离

GPU版本

import numpy as np
import faiss

# 生成一些随机数据
dimension = 128  # 数据的维度
db_size = 10000  # 数据库的大小
np.random.seed(1234)  # 为了结果的可重复性
db_vectors = np.random.random((db_size, dimension)).astype('float32')

# 创建一个FAISS索引
nlist = 100  # 聚类中心的数量
kmeans_index = faiss.IndexFlatL2(dimension)
index = faiss.IndexIVFFlat(kmeans_index, dimension, nlist)

# 将索引转移到GPU上
res = faiss.StandardGpuResources()  # 创建一个GPU资源对象
gpu_index = faiss.index_cpu_to_gpu(res, 0, index)  # 将索引转移到GPU上

# 训练索引
gpu_index.train(db_vectors)

# 将数据添加到索引中
gpu_index.add(db_vectors)

# 保存索引到文件
# 将索引从GPU转移到CPU
cpu_index = faiss.index_gpu_to_cpu(gpu_index)

# 将索引写入文件
faiss.write_index(cpu_index, 'trained_index_file.index')

# 在下次测试的时候，你可以从文件加载索引
# 从文件中读取索引
cpu_index = faiss.read_index('trained_index_file.index')

# 将索引转移到GPU
gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)

# 然后你可以使用加载的索引进行搜索
n_query = 10  # 查询的数量
query_vectors = np.random.random((n_query, dimension)).astype('float32')
k = 4  # 我们想要搜索的最近邻的数量
distances, indices = gpu_index.search(query_vectors, k)

# 打印结果
print(indices)  # 每行包含一个查询的最近邻的索引
print(distances)  # 每行包含一个查询的最近邻的距离