当前位置:   article > 正文

faiss使用实例,快速入门_faiss 例子

faiss 例子

faiss使用实例,快速入门

基本步骤

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

实例

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)  # 每行包含一个查询的最近邻的距离
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26

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

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

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

  1. 生成一些随机数据:这段代码首先生成了一些随机的浮点数向量,这些向量将被用于训练和添加到FAISS索引中。

  2. 创建一个FAISS索引:然后,代码创建了一个FAISS索引。这个索引是IndexIVFFlat类型,这是一种需要训练的索引。

  3. 将索引转移到GPU上:为了加速计算,代码将索引转移到了GPU上。

  4. 训练索引:然后,代码使用生成的随机数据训练了索引。

  5. 将数据添加到索引中:训练后,代码将随机数据添加到了索引中。

  6. 保存索引到文件:然后,代码将索引从GPU转移到CPU,并将其保存到文件中,以便以后使用。

  7. 从文件加载索引:代码从文件中加载了索引,并将其转移到GPU上。

  8. 使用加载的索引进行搜索:最后,代码生成了一些新的随机查询向量,并使用加载的索引对这些查询进行了搜索。搜索的结果是每个查询的最近邻的索引和距离。

总的来说,这段代码展示了如何使用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)  # 每行包含一个查询的最近邻的距离
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47

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

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

在你的代码片段中,_train方法调用了FAISS索引的train方法。这意味着如果索引需要训练(例如,它是一个IndexIVFFlatIndexIVFPQ索引),那么这个方法将执行训练过程。如果索引不需要训练(例如,它是一个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)  # 每行包含一个查询的最近邻的距离
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54

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)  # 每行包含一个查询的最近邻的距离
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/小小林熬夜学编程/article/detail/732525
推荐阅读
相关标签
  

闽ICP备14008679号