Optimizations for template_similarity (numba and dependencies) #3405
Conversation
yger
commented
Sep 13, 2024
yger
commented
- Add the possibility to compute the template similarities (rather slow for the moment) using a numba kernel.
- Remove the dependencies to sklearn in template_similarity without numba, since we do not plan anymore to add metric
- Also adds prange in some other numba functions, with the keyword parallel=True (otherwise, prange is not activated, as far as I understood).
for more information, see https://pre-commit.ci
|
One thing to keep in mind is that the process of creating the threads is pretty expensive, so depending on the number of computations it can actually be beneficial to do the calculation serially rather than do the parallel operation. Not that we shouldn't do this, but we should just thing about common use case vs worst case scenario. |
|
Yes, but here on my machine the speed up is rather drastic, but I can do more benchmarks if needed. In its current form, the template_similarity is too slow (if max_lag_ms is not 0 this is really bad) for large number of templates. So one must optimize it. I've discussed with @samuelgarcia and made two proposals. One in PR #3174 involving PoolExecutor and this one, relying on numba. This is simpler I guess, but this is worth letting @samuelgarcia comment :-) |
| for count, shift in enumerate(shift_loop): | ||
| src_sliced_templates = templates_array[:, num_shifts : num_samples - num_shifts] | ||
| tgt_sliced_templates = other_templates_array[:, num_shifts + shift : num_samples - num_shifts + shift] | ||
| for i in numba.prange(num_templates): |
There was a problem hiding this comment.
I would put the prange the very first loop.
I do not known if the thread are spawn for every shift loop.
In short I would invert the prange on template and the shift loop.
Could you try this and make benchmark ?
There was a problem hiding this comment.
I've tried to invert the loops, and it does not make any difference
for more information, see https://pre-commit.ci
|
For 500 templates with 384 channels (Neuropixels like) we have on a laptop (for 0.1 ms of max_lag_ms)
|
|
@yger could you add a test that checks that numba and numpy implementations give the same results? |
|
Done, tests are passing |
|
Merci Pierre. This is really cool. |
