Use nth_element for median computation in IndexLSH #4653
+8
−6
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Replace std::sort with std::nth_element for median calculation in IndexLSH training.
Contributor
|
@limqiying has imported this pull request. If you are a Meta employee, you can view this in D86234669. |
faiss/IndexLSH.cpp
Outdated
Comment on lines
89
to
98
| // Use nth_element (O(n)) instead of sort (O(n log n)) for median | ||
| if (n % 2 == 1) { | ||
| std::nth_element(xi, xi + n / 2, xi + n); | ||
| thresholds[i] = xi[n / 2]; | ||
| else | ||
| thresholds[i] = (xi[n / 2 - 1] + xi[n / 2]) / 2; | ||
| } else { | ||
| std::nth_element(xi, xi + n / 2, xi + n); | ||
| float median_high = xi[n / 2]; | ||
| std::nth_element(xi, xi + n / 2 - 1, xi + n); | ||
| thresholds[i] = (xi[n / 2 - 1] + median_high) / 2; | ||
| } |
There was a problem hiding this comment.
maybe we can do like this?
// Use nth_element (O(n)) instead of sort (O(n log n)) for median
std::nth_element(xi, xi + n / 2, xi + n);
float median = xi[n / 2];
if (n % 2 == 0) {
std::nth_element(xi, xi + n / 2 - 1, xi + n);
median = (median + xi[n / 2 - 1]) / 2;
}
thresholds[i] = median;
Contributor
Author
There was a problem hiding this comment.
Good point! Your version is cleaner!
|
Hi @Ray0907 , thanks for your contribution. Ditto the other perf changes: did you try benchmarking it? Do you observe speedup? If yes, can you paste the script you used here? |
Contributor
|
I agree that it's the proper way of computing the median but the performance improvement could be minimal. |
| std::nth_element(xi, xi + n / 2, xi + n); | ||
| float median = xi[n / 2]; | ||
| if (n % 2 == 0) { | ||
| std::nth_element(xi, xi + n / 2 - 1, xi + n); |
Contributor
There was a problem hiding this comment.
looks good, unless running 2x std::nth_element() is too costly.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Replace std::sort with std::nth_element for median calculation in
IndexLSH training.