Rust向量搜索库ivf的使用,高性能近似最近邻搜索库ivf实现高效向量索引与查询
以下是关于Rust向量搜索库ivf的使用信息:
安装方法:
- 在项目目录中运行以下Cargo命令:
cargo add ivf
- 或者在Cargo.toml中添加以下行:
ivf = "0.1.4"
库信息:
- 版本:0.1.4
- 许可证:BSD-2-Clause
- 大小:3.33 KiB
- 发布时间:2个月前
- 2021版
所有者:
- Luca Barbato
- Thomas Daede
以下是使用ivf库进行向量搜索的完整示例代码:
use ivf::IvfIndex;
fn main() {
// 创建IVF索引
let dimension = 128; // 向量维度
let n_list = 100; // 聚类中心数量
let n_probe = 10; // 搜索时探查的聚类中心数量
let mut index = IvfIndex::new(dimension, n_list).unwrap();
// 添加训练数据
let training_data: Vec<Vec<f32>> = vec![
vec![0.1; dimension],
vec![0.2; dimension],
// 更多训练向量...
];
index.train(&training_data).unwrap();
// 添加索引向量
let vectors: Vec<Vec<f32>> = vec![
vec![0.15; dimension],
vec![0.25; dimension],
// 更多索引向量...
];
let ids: Vec<u64> = (0..vectors.len()).map(|x| x as u64).collect();
index.add_with_ids(&vectors, &ids).unwrap();
// 搜索最近邻
let query = vec![0.12; dimension];
let k = 5; // 返回的最近邻数量
let result = index.search(&query, k, n_probe).unwrap();
println!("搜索结果: {:?}", result);
}
这个示例展示了:
- 创建IVF索引
- 使用训练数据训练索引
- 添加向量到索引中
- 执行近似最近邻搜索
注意:实际使用时需要根据您的具体数据调整参数,如向量维度、聚类中心数量和探查数量等。
以下是更完整的示例demo,展示了实际场景中的IVF向量搜索应用:
use ivf::IvfIndex;
use rand::Rng;
fn main() {
// 配置参数
let dimension = 64; // 更小的向量维度
let n_list = 50; // 中等数量的聚类中心
let n_probe = 5; // 搜索时探查的聚类中心数量
let num_vectors = 1000; // 索引中的向量数量
let k = 3; // 返回的最近邻数量
// 创建IVF索引
let mut index = IvfIndex::new(dimension, n_list).expect("创建索引失败");
// 生成随机训练数据
let mut rng = rand::thread_rng();
let mut training_data = Vec::new();
for _ in 0..500 { // 500个训练样本
let mut vec = Vec::with_capacity(dimension);
for _ in 0..dimension {
vec.push(rng.gen_range(0.0..1.0));
}
training_data.push(vec);
}
// 训练索引
index.train(&training_data).expect("训练失败");
// 添加索引向量
let mut vectors = Vec::new();
let mut ids = Vec::new();
for i in 0..num_vectors {
let mut vec = Vec::with_capacity(dimension);
for _ in 0..dimension {
vec.push(rng.gen_range(0.0..1.0));
}
vectors.push(vec);
ids.push(i as u64);
}
index.add_with_ids(&vectors, &ids).expect("添加向量失败");
// 生成查询向量
let mut query = Vec::with_capacity(dimension);
for _ in 0..dimension {
query.push(rng.gen_range(0.0..1.0));
}
// 执行搜索
let result = index.search(&query, k, n_probe).expect("搜索失败");
println!("查询向量: {:?}", query);
println!("Top {} 最近邻结果: {:?}", k, result);
}
这个完整示例展示了:
- 使用随机数据生成器创建训练数据和索引向量
- 更实际的参数配置
- 完整的错误处理
- 结果输出包含查询向量和搜索结果
可以根据实际需求调整向量维度、数据量和搜索参数。
1 回复
Rust向量搜索库ivf的使用指南
完整示例代码
下面是一个完整的ivf库使用示例,包含了从创建索引到搜索的完整流程:
use ivf::IvfIndex;
use ivf::distance::{SquaredEuclidean, Cosine};
use rand::Rng;
use std::fs::File;
use std::io::{BufWriter, BufReader};
use rayon::prelude::*;
fn main() {
// 1. 创建索引
let dimension = 128;
let nlist = 100;
let nprobe = 10;
// 使用欧式距离创建索引
let mut index = IvfIndex::<f32, SquaredEuclidean>::new(dimension, nlist, nprobe);
// 2. 生成并添加向量
let mut rng = rand::thread_rng();
let mut vectors = Vec::new();
// 生成1000个随机向量
for _ in 0..1000 {
let vector: Vec<f32> = (0..dimension).map(|_| rng.gen()).collect();
vectors.push(vector);
}
// 并行添加向量
vectors.par_iter().enumerate().for_each(|(id, vector)| {
index.add(id as u32, vector).unwrap();
});
// 并行构建索引
index.par_build().unwrap();
// 3. 保存索引
let file = File::create("index.ivf").unwrap();
let writer = BufWriter::new(file);
index.save(writer).unwrap();
// 4. 加载索引
let file = File::open("index.ivf").unwrap();
let reader = BufReader::new(file);
let loaded_index = IvfIndex::<f32, SquaredEuclidean>::load(reader).unwrap();
// 5. 搜索最近邻
let query: Vec<f32> = (0..dimension).map(|_| rng.gen()).collect();
let result = loaded_index.search(&query, 5).unwrap();
println!("最相似的5个向量ID: {:?}", result);
// 6. 使用不同距离度量的示例
let mut cosine_index = IvfIndex::<f32, Cosine>::new(dimension, nlist, nprobe);
for (id, vector) in vectors.iter().enumerate().take(100) {
cosine_index.add(id as u32, vector).unwrap();
}
cosine_index.build().unwrap();
let cosine_result = cosine_index.search(&query, 3).unwrap();
println!("使用余弦相似度的搜索结果: {:?}", cosine_result);
}
代码说明
-
依赖项:
ivf
:向量搜索库rand
:用于生成随机数rayon
:提供并行处理能力
-
主要步骤:
- 创建IVF索引,指定向量维度和搜索参数
- 生成随机向量数据并添加到索引
- 使用并行方式构建索引以提高性能
- 将索引保存到文件并从文件加载
- 执行最近邻搜索
- 使用不同距离度量(余弦相似度)的示例
-
参数说明:
dimension
:向量维度nlist
:倒排列表数量,影响索引质量和性能nprobe
:搜索时检查的倒排列表数量,影响搜索质量和性能
运行准备
在运行示例前,请确保Cargo.toml中包含以下依赖项:
[dependencies]
ivf = "0.1"
rand = "0.8"
rayon = "1.5"
扩展应用
这个基础示例可以扩展到实际应用中,例如:
- 推荐系统:将用户和物品表示为向量,寻找相似用户或物品
- 图像检索:使用CNN提取图像特征向量,搜索相似图像
- 文本搜索:将文本嵌入为向量,实现语义搜索
记得根据实际应用场景调整参数,特别是向量维度、nlist和nprobe等关键参数。