iOS 中的矢量相似度搜索

问题描述 投票:0回答:1

是否有适用于 iOS 的矢量相似性搜索的实现?

我有一组 ~10K+ 向量。当我得到一个新向量时,我想从集合中找到与新向量最相似的前 K 个向量。

这可以在循环中完成:计算集合中每个向量与新向量之间的相似度,然后选择顶部 K。

有更有效的方法(例如 FAISSScanNN)可以扩展到数百万个向量,但它们不适用于 iOS。是否有适用于 iOS 的矢量相似性搜索的实现?

编辑1: 我找到了这个项目https://github.com/hora-search/hora-ios,但它看起来已经死了。有人试过吗?

ios swift vector cosine-similarity faiss
1个回答
0
投票

想法

对于我来说,如果要查询的向量大小在10k+左右,我想使用暴力查询的方式。 实际上,在 iOS 加速框架中,有 cblas_ 函数(目前可能在 API 参考中标记为已弃用)。 由于 faiss 中的 FlatIndex 实际上是做 Gemm 然后 TopK selection

我建议可以从 cblas_sgemm + heap sort(topk) 开始。在我看来,在移动平台上,省电可能更重要,所以更喜欢调用厂商的原生API。

一个例子


// return random vector with L2 norm
func randFloatArray(dims: Int, count: Int) -> [Float]
{
    var values = [Float](repeating: 0.0, count: dims * count)
    for n in 0...count-1
    {
        var sum: Float = 0.0
        let offset = n * dims
        for index in 0...dims-1
        {
            let val = Float(drand48());
            sum += val * val
            values[offset+index] = val
        }
        let sum_sqrt = sqrt(sum)
        for index in 0...dims-1
        {
            values[offset+index] /= sum_sqrt
        }
    }
    return values
}



class Vectors {
    struct QueryResult {
        var v: Float
        var idx: Int64
    }
    var data: [Float]
    var ids: [Int64]
    var dims: Int32
    init (dims: Int) {
        self.data = [Float]()
        self.ids = [Int64]()
        self.dims = Int32(dims)
    }
    
    // heap sort for select top k
    func heapSort(_ array: [QueryResult], topk: Int) -> [QueryResult] {
        var heap = Array(array.prefix(topk))

        func siftDown(_ start: Int, _ end: Int) {
            var root = start
            while root * 2 + 1 <= end {
                let child = root * 2 + 1
                var swap = root
                if heap[swap].v > heap[child].v {
                    swap = child
                }
                if child + 1 <= end && heap[swap].v > heap[child + 1].v {
                    swap = child + 1
                }
                if swap == root {
                    return
                } else {
                    heap.swapAt(root, swap)
                    root = swap
                }
            }
        }

        let count = heap.count
        for i in stride(from: count / 2 - 1, through: 0, by: -1) {
            siftDown(i, count - 1)
        }

        for i in stride(from: count, to: array.count, by: 1) {
            if array[i].v > heap[0].v {
                heap[0] = array[i]
                siftDown(0, count - 1)
            }
        }
        return heap.sorted(by:{$0.v > $1.v})
    }
    
    func insert(vec: [Float], ids: [Int64])
    {
        self.data.append(contentsOf: vec)
        self.ids.append(contentsOf: ids)
    }
    
    func count() -> Int
    {
        return data.count / Int(dims)
    }
    
    // query, return (results, gemm time, k select time)
    func query(vec: [Float], topk: Int) -> ([QueryResult], UInt64, UInt64)
    {
        var c = [Float](repeating: 0.0, count: Int(self.count()))
        var results = [QueryResult](repeating: QueryResult(v: 0.0, idx: -1), count: Int(self.count()))

        let t0 = UInt64(Date().timeIntervalSince1970 * 1000)
        
        vec.withUnsafeBufferPointer { vec in
            data.withUnsafeBufferPointer {data in
                c.withUnsafeMutableBufferPointer {c in
                    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, 1, Int32(self.count()), dims, 1.0, vec.baseAddress, dims, data.baseAddress, dims, 0.0, c.baseAddress, Int32(self.count()))
                }
            }
        }

        let t1 = UInt64(Date().timeIntervalSince1970 * 1000)
        // prepare result before sort
        for index in 0...self.count()-1 {
            results[index] = QueryResult(v:c[index], idx: self.ids[index])
        }
        // sort for poc, using heap sort for tokk is more efficient
        results = heapSort(results, topk: topk)
        let t2 = UInt64(Date().timeIntervalSince1970 * 1000)
        return (results, t1-t0, t2-t1)
    }
}


iOS 设备上的一些基准测试

256维,输入1个向量查询100K条记录中的top 3。

© www.soinside.com 2019 - 2024. All rights reserved.