diff --git a/contrib/exhaustive_search.py b/contrib/exhaustive_search.py index b7b57c9ea..fa7f1b11c 100644 --- a/contrib/exhaustive_search.py +++ b/contrib/exhaustive_search.py @@ -20,7 +20,7 @@ def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2): LOG.info("knn_ground_truth queries size %s k=%d" % (xq.shape, k)) t0 = time.time() nq, d = xq.shape - keep_max = metric_type == faiss.METRIC_INNER_PRODUCT + keep_max = faiss.is_similarity_metric(metric_type) rh = faiss.ResultHeap(nq, k, keep_max=keep_max) index = faiss.IndexFlat(d, metric_type) @@ -265,7 +265,7 @@ def range_search_max_results(index, query_iterator, radius, (totres, max_results)) radius, totres = apply_maxres( res_batches, min_results, - keep_max=index.metric_type == faiss.METRIC_INNER_PRODUCT + keep_max=faiss.is_similarity_metric(index.metric_type) ) t2 = time.time() t_search += t1 - t0 @@ -281,7 +281,7 @@ def range_search_max_results(index, query_iterator, radius, if clip_to_min and totres > min_results: radius, totres = apply_maxres( res_batches, min_results, - keep_max=index.metric_type == faiss.METRIC_INNER_PRODUCT + keep_max=faiss.is_similarity_metric(index.metric_type) ) nres = np.hstack([nres_i for nres_i, dis_i, ids_i in res_batches]) diff --git a/contrib/ivf_tools.py b/contrib/ivf_tools.py index 581e55f4f..9562240f7 100644 --- a/contrib/ivf_tools.py +++ b/contrib/ivf_tools.py @@ -138,7 +138,7 @@ class BigBatchSearcher: self.xq = xq self.index = index self.use_float16 = use_float16 - keep_max = index.metric_type == faiss.METRIC_INNER_PRODUCT + keep_max = faiss.is_similarity_metric(index.metric_type) self.rh = faiss.ResultHeap(len(xq), k, keep_max=keep_max) self.t_accu = [0] * 5 self.t_display = self.t0 = time.time() diff --git a/faiss/Clustering.cpp b/faiss/Clustering.cpp index 0a0249f89..7cf2ba094 100644 --- a/faiss/Clustering.cpp +++ b/faiss/Clustering.cpp @@ -372,7 +372,7 @@ void Clustering::train_encoded( std::unique_ptr dis(new float[nx]); // remember best iteration for redo - bool lower_is_better = index.metric_type != METRIC_INNER_PRODUCT; + bool lower_is_better = !is_similarity_metric(index.metric_type); float best_obj = lower_is_better ? HUGE_VALF : -HUGE_VALF; std::vector best_iteration_stats; std::vector best_centroids; diff --git a/faiss/IndexFlat.cpp b/faiss/IndexFlat.cpp index 5a8ebd82f..0fa3b8206 100644 --- a/faiss/IndexFlat.cpp +++ b/faiss/IndexFlat.cpp @@ -40,6 +40,10 @@ void IndexFlat::search( } else if (metric_type == METRIC_L2) { float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances}; knn_L2sqr(x, get_xb(), d, n, ntotal, &res, nullptr, sel); + } else if (is_similarity_metric(metric_type)) { + float_minheap_array_t res = {size_t(n), size_t(k), labels, distances}; + knn_extra_metrics( + x, get_xb(), d, n, ntotal, metric_type, metric_arg, &res); } else { FAISS_THROW_IF_NOT(!sel); float_maxheap_array_t res = {size_t(n), size_t(k), labels, distances}; diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp index 4c6757480..f133baf64 100644 --- a/faiss/IndexHNSW.cpp +++ b/faiss/IndexHNSW.cpp @@ -98,7 +98,7 @@ struct NegativeDistanceComputer : DistanceComputer { }; DistanceComputer* storage_distance_computer(const Index* storage) { - if (storage->metric_type == METRIC_INNER_PRODUCT) { + if (is_similarity_metric(storage->metric_type)) { return new NegativeDistanceComputer(storage->get_distance_computer()); } else { return storage->get_distance_computer(); @@ -346,7 +346,7 @@ void IndexHNSW::search( InterruptCallback::check(); } - if (metric_type == METRIC_INNER_PRODUCT) { + if (is_similarity_metric(metric_type)) { // we need to revert the negated distances for (size_t i = 0; i < k * n; i++) { distances[i] = -distances[i]; diff --git a/faiss/IndexIVFAdditiveQuantizer.cpp b/faiss/IndexIVFAdditiveQuantizer.cpp index 7ba8ddb26..0fa836aa0 100644 --- a/faiss/IndexIVFAdditiveQuantizer.cpp +++ b/faiss/IndexIVFAdditiveQuantizer.cpp @@ -145,7 +145,7 @@ struct AQInvertedListScanner : InvertedListScanner { : ia(ia), aq(*ia.aq) { this->store_pairs = store_pairs; this->code_size = ia.code_size; - keep_max = ia.metric_type == METRIC_INNER_PRODUCT; + keep_max = is_similarity_metric(ia.metric_type); tmp.resize(ia.d); } diff --git a/faiss/IndexNNDescent.cpp b/faiss/IndexNNDescent.cpp index ec7bf188c..8cdc0c4ab 100644 --- a/faiss/IndexNNDescent.cpp +++ b/faiss/IndexNNDescent.cpp @@ -88,7 +88,7 @@ struct NegativeDistanceComputer : DistanceComputer { }; DistanceComputer* storage_distance_computer(const Index* storage) { - if (storage->metric_type == METRIC_INNER_PRODUCT) { + if (is_similarity_metric(storage->metric_type)) { return new NegativeDistanceComputer(storage->get_distance_computer()); } else { return storage->get_distance_computer(); diff --git a/faiss/IndexNSG.cpp b/faiss/IndexNSG.cpp index d2fe34401..a7cfd490a 100644 --- a/faiss/IndexNSG.cpp +++ b/faiss/IndexNSG.cpp @@ -112,7 +112,7 @@ void IndexNSG::search( InterruptCallback::check(); } - if (metric_type == METRIC_INNER_PRODUCT) { + if (is_similarity_metric(metric_type)) { // we need to revert the negated distances for (size_t i = 0; i < k * n; i++) { distances[i] = -distances[i]; diff --git a/faiss/MetaIndexes.cpp b/faiss/MetaIndexes.cpp index 3499a5642..50c7c88cb 100644 --- a/faiss/MetaIndexes.cpp +++ b/faiss/MetaIndexes.cpp @@ -218,7 +218,7 @@ void IndexRandom::search( float dprev = 0; for (int j = 0; j < k; j++) { float step = rng.rand_float(); - if (metric_type == METRIC_INNER_PRODUCT) { + if (is_similarity_metric(metric_type)) { step = -step; } dprev += step; diff --git a/faiss/MetricType.h b/faiss/MetricType.h index 404028145..538b0a8e7 100644 --- a/faiss/MetricType.h +++ b/faiss/MetricType.h @@ -31,11 +31,20 @@ enum MetricType { METRIC_Canberra = 20, METRIC_BrayCurtis, METRIC_JensenShannon, + METRIC_Jaccard, ///< defined as: sum_i(min(a_i, b_i)) / sum_i(max(a_i, b_i)) + ///< where a_i, b_i > 0 }; /// all vector indices are this type using idx_t = int64_t; +/// this function is used to distinguish between min and max indexes since +/// we need to support similarity and dis-similarity metrics in a flexible way +constexpr bool is_similarity_metric(MetricType metric_type) { + return ((metric_type == METRIC_INNER_PRODUCT) || + (metric_type == METRIC_Jaccard)); +} + } // namespace faiss #endif diff --git a/faiss/gpu/impl/Distance.cuh b/faiss/gpu/impl/Distance.cuh index 6c1a3e019..b9ea4242f 100644 --- a/faiss/gpu/impl/Distance.cuh +++ b/faiss/gpu/impl/Distance.cuh @@ -228,6 +228,13 @@ void allPairwiseDistanceOnDevice( outDistances, JensenShannonDistance(), stream); + } else if (metric == faiss::MetricType::METRIC_Jaccard) { + runGeneralDistanceKernel( + tVectorsDimInnermost, + tQueriesDimInnermost, + outDistances, + JaccardSimilarity(), + stream); } else { FAISS_THROW_FMT("unimplemented metric type %d", metric); } @@ -375,6 +382,16 @@ void bfKnnOnDevice( JensenShannonDistance(), outDistances, outIndices); + } else if (metric == faiss::MetricType::METRIC_Jaccard) { + runGeneralDistance( + resources, + stream, + tVectorsDimInnermost, + tQueriesDimInnermost, + k, + JaccardSimilarity(), + outDistances, + outIndices); } else { FAISS_THROW_FMT("unimplemented metric type %d", metric); } diff --git a/faiss/gpu/impl/DistanceUtils.cuh b/faiss/gpu/impl/DistanceUtils.cuh index 56c6c88bc..530382b40 100644 --- a/faiss/gpu/impl/DistanceUtils.cuh +++ b/faiss/gpu/impl/DistanceUtils.cuh @@ -251,6 +251,35 @@ struct JensenShannonDistance { float dist; }; +struct JaccardSimilarity { + __host__ __device__ JaccardSimilarity() : numerator(0), denominator(0) {} + + static constexpr bool kDirection = true; // maximize + static constexpr float kIdentityData = 0; + static constexpr float kMaxDistance = -std::numeric_limits::max(); + + __host__ __device__ void handle(float a, float b) { + numerator += fmin(a, b); + denominator += fmax(a, b); + } + + __host__ __device__ float reduce() { + return numerator / denominator; + } + + __host__ __device__ void combine(const JaccardSimilarity& v) { + numerator += v.numerator; + denominator += v.denominator; + } + + __host__ __device__ JaccardSimilarity zero() const { + return JaccardSimilarity(); + } + + float numerator; + float denominator; +}; + template Tensor sliceCentroids( Tensor& centroids, diff --git a/faiss/gpu/test/TestGpuDistance.cu b/faiss/gpu/test/TestGpuDistance.cu index cb9b1a417..eedec6b70 100644 --- a/faiss/gpu/test/TestGpuDistance.cu +++ b/faiss/gpu/test/TestGpuDistance.cu @@ -37,7 +37,8 @@ void testTransposition( std::vector vecs = randVecs(numVecs, dim); std::vector queries = randVecs(numQuery, dim); - if (metric == faiss::MetricType::METRIC_JensenShannon) { + if ((metric == faiss::MetricType::METRIC_JensenShannon) || + (metric == faiss::MetricType::METRIC_Jaccard)) { // make values positive for (auto& v : vecs) { v = std::abs(v); @@ -192,6 +193,10 @@ TEST(TestGpuDistance, JensenShannon) { testTransposition(false, false, faiss::MetricType::METRIC_JensenShannon); } +TEST(TestGpuDistance, Jaccard) { + testTransposition(false, false, faiss::MetricType::METRIC_Jaccard); +} + int main(int argc, char** argv) { testing::InitGoogleTest(&argc, argv); diff --git a/faiss/gpu/test/test_gpu_basics.py b/faiss/gpu/test/test_gpu_basics.py index afc12080b..69633d4d6 100755 --- a/faiss/gpu/test/test_gpu_basics.py +++ b/faiss/gpu/test/test_gpu_basics.py @@ -305,7 +305,8 @@ class TestAllPairwiseDistance(unittest.TestCase): faiss.METRIC_Linf, faiss.METRIC_Canberra, faiss.METRIC_BrayCurtis, - faiss.METRIC_JensenShannon + faiss.METRIC_JensenShannon, + faiss.METRIC_Jaccard ] for metric in metrics: @@ -349,7 +350,7 @@ class TestAllPairwiseDistance(unittest.TestCase): # INNER_PRODUCT is in descending order, make sure it is the same # order - if metric == faiss.METRIC_INNER_PRODUCT: + if faiss.is_similarity_metric(metric): ref_d = np.sort(ref_d, axis=1) print('f32', np.abs(ref_d - out_d).max()) @@ -382,7 +383,7 @@ class TestAllPairwiseDistance(unittest.TestCase): # INNER_PRODUCT is in descending order, make sure it is the same # order - if metric == faiss.METRIC_INNER_PRODUCT: + if faiss.is_similarity_metric(metric): ref_d_f16 = np.sort(ref_d_f16, axis=1) print('f16', np.abs(ref_d_f16 - out_d_f16).max()) diff --git a/faiss/impl/NSG.cpp b/faiss/impl/NSG.cpp index cc4ab84ad..1f30b576b 100644 --- a/faiss/impl/NSG.cpp +++ b/faiss/impl/NSG.cpp @@ -57,7 +57,7 @@ struct NegativeDistanceComputer : DistanceComputer { } // namespace DistanceComputer* storage_distance_computer(const Index* storage) { - if (storage->metric_type == METRIC_INNER_PRODUCT) { + if (is_similarity_metric(storage->metric_type)) { return new NegativeDistanceComputer(storage->get_distance_computer()); } else { return storage->get_distance_computer(); diff --git a/faiss/utils/extra_distances-inl.h b/faiss/utils/extra_distances-inl.h index 0463385d2..d3768df66 100644 --- a/faiss/utils/extra_distances-inl.h +++ b/faiss/utils/extra_distances-inl.h @@ -8,6 +8,7 @@ /** In this file are the implementations of extra metrics beyond L2 * and inner product */ +#include #include #include @@ -17,12 +18,13 @@ template struct VectorDistance { size_t d; float metric_arg; + static constexpr bool is_similarity = is_similarity_metric(mt); inline float operator()(const float* x, const float* y) const; // heap template to use for this type of metric using C = typename std::conditional< - mt == METRIC_INNER_PRODUCT, + is_similarity_metric(mt), CMin, CMax>::type; }; @@ -114,4 +116,18 @@ inline float VectorDistance::operator()( return 0.5 * accu; } +template <> +inline float VectorDistance::operator()( + const float* x, + const float* y) const { + // WARNING: this distance is defined only for positive input vectors. + // Providing vectors with negative values would lead to incorrect results. + float accu_num = 0, accu_den = 0; + for (size_t i = 0; i < d; i++) { + accu_num += fmin(x[i], y[i]); + accu_den += fmax(x[i], y[i]); + } + return accu_num / accu_den; +} + } // namespace faiss diff --git a/faiss/utils/extra_distances.cpp b/faiss/utils/extra_distances.cpp index dc60611a9..8c0699880 100644 --- a/faiss/utils/extra_distances.cpp +++ b/faiss/utils/extra_distances.cpp @@ -50,14 +50,14 @@ void pairwise_extra_distances_template( } } -template +template void knn_extra_metrics_template( VD vd, const float* x, const float* y, size_t nx, size_t ny, - float_maxheap_array_t* res) { + HeapArray* res) { size_t k = res->k; size_t d = vd.d; size_t check_period = InterruptCallback::get_period_hint(ny * d); @@ -74,16 +74,21 @@ void knn_extra_metrics_template( float* simi = res->get_val(i); int64_t* idxi = res->get_ids(i); - maxheap_heapify(k, simi, idxi); + // maxheap_heapify(k, simi, idxi); + heap_heapify(k, simi, idxi); for (j = 0; j < ny; j++) { float disij = vd(x_i, y_j); - if (disij < simi[0]) { - maxheap_replace_top(k, simi, idxi, disij, j); + // if (disij < simi[0]) { + if ((!vd.is_similarity && (disij < simi[0])) || + (vd.is_similarity && (disij > simi[0]))) { + // maxheap_replace_top(k, simi, idxi, disij, j); + heap_replace_top(k, simi, idxi, disij, j); } y_j += d; } - maxheap_reorder(k, simi, idxi); + // maxheap_reorder(k, simi, idxi); + heap_reorder(k, simi, idxi); } InterruptCallback::check(); } @@ -158,12 +163,14 @@ void pairwise_extra_distances( HANDLE_VAR(BrayCurtis); HANDLE_VAR(JensenShannon); HANDLE_VAR(Lp); + HANDLE_VAR(Jaccard); #undef HANDLE_VAR default: FAISS_THROW_MSG("metric type not implemented"); } } +template void knn_extra_metrics( const float* x, const float* y, @@ -172,7 +179,7 @@ void knn_extra_metrics( size_t ny, MetricType mt, float metric_arg, - float_maxheap_array_t* res) { + HeapArray* res) { switch (mt) { #define HANDLE_VAR(kw) \ case METRIC_##kw: { \ @@ -187,12 +194,33 @@ void knn_extra_metrics( HANDLE_VAR(BrayCurtis); HANDLE_VAR(JensenShannon); HANDLE_VAR(Lp); + HANDLE_VAR(Jaccard); #undef HANDLE_VAR default: FAISS_THROW_MSG("metric type not implemented"); } } +template void knn_extra_metrics>( + const float* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + MetricType mt, + float metric_arg, + HeapArray>* res); + +template void knn_extra_metrics>( + const float* x, + const float* y, + size_t d, + size_t nx, + size_t ny, + MetricType mt, + float metric_arg, + HeapArray>* res); + FlatCodesDistanceComputer* get_extra_distance_computer( size_t d, MetricType mt, @@ -213,6 +241,7 @@ FlatCodesDistanceComputer* get_extra_distance_computer( HANDLE_VAR(BrayCurtis); HANDLE_VAR(JensenShannon); HANDLE_VAR(Lp); + HANDLE_VAR(Jaccard); #undef HANDLE_VAR default: FAISS_THROW_MSG("metric type not implemented"); diff --git a/faiss/utils/extra_distances.h b/faiss/utils/extra_distances.h index 17cd25e35..79b65bc1e 100644 --- a/faiss/utils/extra_distances.h +++ b/faiss/utils/extra_distances.h @@ -33,6 +33,7 @@ void pairwise_extra_distances( int64_t ldb = -1, int64_t ldd = -1); +template void knn_extra_metrics( const float* x, const float* y, @@ -41,7 +42,7 @@ void knn_extra_metrics( size_t ny, MetricType mt, float metric_arg, - float_maxheap_array_t* res); + HeapArray* res); /** get a DistanceComputer that refers to this type of distance and * indexes a flat array of size nb */ diff --git a/tests/test_extra_distances.py b/tests/test_extra_distances.py index eec955e14..a474dd6ba 100644 --- a/tests/test_extra_distances.py +++ b/tests/test_extra_distances.py @@ -82,6 +82,18 @@ class TestExtraDistances(unittest.TestCase): self.run_simple_dis_test(scipy.spatial.distance.jensenshannon, faiss.METRIC_JensenShannon) + def test_jaccard(self): + xq, yb = self.make_example() + ref_dis = np.array([ + [ + (np.min([x, y], axis=0).sum() / np.max([x, y], axis=0).sum()) + for y in yb + ] + for x in xq + ]) + new_dis = faiss.pairwise_distances(xq, yb, faiss.METRIC_Jaccard) + self.assertTrue(np.allclose(ref_dis, new_dis)) + class TestKNN(unittest.TestCase): """ test that the knn search gives the same as distance matrix + argmin """