如何将 numpy 数组保存在 .fvecs 中

Question

我正在使用 faiss 创建向量嵌入的索引，然后可以将其用于执行相似性搜索。由于我的数据集达到数十亿，我正在寻找极快（接近最佳）的解决方案。我发现 faiss 对于不需要将整个数据集加载到内存中进行索引的情况有很好的支持。然而，所有代码演示和示例（基准脚本）都使用 .fvecs 文件，我认为这是一个很好的方法，因为它们可以映射（通过 numpy）。但我不确定，如何将作为火炬张量返回的嵌入存储为 .fvecs 格式？

我当前的实现是将

vecs

保存为 .fvecs 文件：

def write_fvec(filename, vecs):
    with open(filename, "wb") as f:
        nvecs, dim = vecs.shape
        f.write(struct.pack('<i', nvecs))
        f.write(struct.pack('<i', dim))
        vecs.astype('float32').flatten().tofile(f)

a = np.random.rand((1000000, 1024))
write_fvec("test_file.fvecs", a)

但是当我使用它来使用 faiss' code 来读取它时，它使用给定的函数来读取它，

def mmap_fvecs(fname):
    x = np.memmap(fname, dtype='int32', mode='r')
    d = x[0]
    return x.view('float32').reshape(-1, d + 1)[:, 1:]

我收到以下错误：

ValueError: cannot reshape array of size 1024000002 into shape (1000001)

我正在使用维度数组

1000000x1024

Answer 1

有趣的问题，为了确保正确创建和读取 .fvecs 文件，我们将回顾写入和读取这些文件的过程。

写入 .fvecs 文件 .fvecs格式一般存储数据如下：

1. 表示向量数量的整数。
2. 表示每个向量维度的整数。
3. 向量本身，存储为 float32。

写入文件时，请确保每个向量都以 float32 格式正确写入。您提供的编写代码似乎是正确的：

import numpy as np
import struct

def write_fvec(filename, vecs):
    with open(filename, "wb") as f:
        nvecs, dim = vecs.shape
        f.write(struct.pack('<i', nvecs))
        f.write(struct.pack('<i', dim))
        vecs.astype('float32').flatten().tofile(f)

# Generate sample data
a = np.random.rand(1000000, 1024)
write_fvec("test_file.fvecs", a)

读取.fvecs文件 要读取该文件，重要的是要考虑前两个整数 (int32) 不应被解释为 float32 数据的一部分。您的读取函数需要进行调整以处理此问题：

def mmap_fvecs(fname):
    # offset of 8 to skip the first two integers
    x = np.memmap(fname, dtype='float32', mode='r', offset=8)

    # Read nvecs and dim of the first 8 bytes  
    nvecs, dim = struct.unpack('<ii', x[:8].view('int32'))  
    return x.reshape(-1, dim)

# Use the function to read
vecs = mmap_fvecs("test_file.fvecs")
print(vecs.shape)

需要考虑的问题：

确保 np.memmap 中的偏移量正确：偏移量必须为 8 个字节才能跳过前两个整数 (

nvecs and dim

)，每个整数为 4 个字节。

验证读取数组的形状是否符合您的期望： 上述错误表明重塑未正确完成，可能是由于读取文件时对尺寸和尺寸的处理不正确。

希望对你有帮助