From 78d2a8259d5bc97e42bd25793ec878d582e42b2e Mon Sep 17 00:00:00 2001 From: zhenghaoz Date: Tue, 1 Feb 2022 20:28:55 +0800 Subject: [PATCH] evaluate HNSW in parallel (#371) --- base/search/hnsw.go | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/base/search/hnsw.go b/base/search/hnsw.go index f3c8c058d..2a3580c39 100644 --- a/base/search/hnsw.go +++ b/base/search/hnsw.go @@ -332,18 +332,23 @@ func recall(expected, actual []int32) float32 { return result / float32(len(actual)) } -func (b *HNSWBuilder) evaluate(idx VectorIndex, prune0 bool) float32 { +func (b *HNSWBuilder) evaluate(idx *HNSW, prune0 bool) float32 { testSize := mathutil.Min(b.testSize, len(b.data)) samples := b.rng.Sample(0, len(b.data), testSize) var result, count float32 - for _, i := range samples { - expected, _ := b.bruteForce.Search(b.data[i], b.k, prune0) + var mu sync.Mutex + _ = base.Parallel(len(samples), idx.numJobs, func(_, i int) error { + sample := samples[i] + expected, _ := b.bruteForce.Search(b.data[sample], b.k, prune0) if len(expected) > 0 { - actual, _ := idx.Search(b.data[i], b.k, prune0) + actual, _ := idx.Search(b.data[sample], b.k, prune0) + mu.Lock() + defer mu.Unlock() result += recall(expected, actual) count++ } - } + return nil + }) return result / count }