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 :-) |
| @@ -147,10 +154,144 @@ def _get_data(self): | |||
| compute_template_similarity = ComputeTemplateSimilarity.function_factory() | |||
|
|
|||
|
|
|||
| if HAVE_NUMBA: | |||
There was a problem hiding this comment.
When we were making the import faster I remember that we discussed whether this compiles at init time or not. Does anyone have a better memory than me @zm711 @JoeZiminski ?
There was a problem hiding this comment.
I think the compiling at init is mainly due to type hints. I was actually listening to someone talk about numba the other day and they recommend not to do the type hints for numba and let it infer them. So I think we should be fine. We should see the slow down in the import action right? or we could use tuna (I forget the name @JoeZiminski )
There was a problem hiding this comment.
Ah yes, it was because Aurelien specified the actual types on the numba decorator. Now I remember.
| src_template = src_sliced_templates[i] | ||
| overlapping_templates = np.flatnonzero(np.sum(mask[i], 1)) | ||
| tgt_templates = tgt_sliced_templates[overlapping_templates] | ||
| for gcount, j in enumerate(overlapping_templates): |
There was a problem hiding this comment.
not sure if enumerate is good for numba or not.
There was a problem hiding this comment.
https://numba.pydata.org/numba-doc/0.17.0/reference/pysupported.html
it is supported. Not sure if it is a good idea.
| 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 ?
| norm_j = 0 | ||
|
|
||
| for k in range(len(src)): | ||
| if method == "l1": |
There was a problem hiding this comment.
Maybe we should use constant with number to avoid testing against a str object in any inner loops.
Something in C style like
L1 = 0
L2 = 1
COSINE = 2
method_int = {"l1": L1, "l2": L2, "cosine": COSINE}[method]
for ... in ... :
for ... in ... :
for ... in ... :
if method_int == L1:
| src = src_template[:, mask[i, j]].flatten() | ||
| tgt = (tgt_templates[gcount][:, mask[i, j]]).flatten() |
There was a problem hiding this comment.
Cool but how does numba deal with method ? Is it fast ?
There was a problem hiding this comment.
https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html
it is supported, for "c" order only. Which I think we always follow? Again not sure on speed.
| if method == "l1": | ||
| distances[count, i, j] /= norm_i + norm_j | ||
| elif method == "l2": | ||
| norm_i = np.sqrt(norm_i) |
There was a problem hiding this comment.
Do we need np.sqrt or sqrt in numba ?
