Add array api support for jaccard score #31204
Merged
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
|
Some benchmarks data size = 1e7
|
|
CC: @ogrisel @betatim @adrinjalali for reviews. It seems like Pytorch CPU degrades the performance. |
virchan
approved these changes
Apr 21, 2025
There was a problem hiding this comment.
Thanks for the PR, @OmarManzoor!
Just one small nitpick from my side—otherwise, LGTM!
Edit: I forget we haven't checked with MPS. 😅
adrinjalali
approved these changes
Apr 28, 2025
|
I guess we're okay that PyTorch CPU has a subpar implementation? |
I think so. Just to be clear I tested on a kaggle kernel which sometimes involves differing and conflicting versions of packages but I think that is the best we can get for publicly available free gpus. Since metrics do not involve too much computation using Pytorch with CPU doesn't really offer much benefit and we are better off with the original numpy implementation when we are using a CPU. But I think let's get an opinion from @ogrisel as well. |
|
@ogrisel Do you think this can be merged? |
|
Let's merge this one, thanks! I had a closer look at the performance difference between numpy and PyTorch CPU. Looking a bit at it, this is due to at least 2 things:
Quick benchmark code import numpy as np
import torch
from sklearn.metrics import jaccard_score
import sklearn
y_true = np.ones((10_000, 1000), dtype=np.int64)
y_pred = np.ones((10_000, 1000), dtype=np.int64)
y_true_torch = torch.asarray(y_true)
y_pred_torch = torch.asarray(y_pred)
print('numpy')
%timeit jaccard_score(y_true, y_pred, average='macro')
print('torch')
with sklearn.config_context(array_api_dispatch=True):
%timeit jaccard_score(y_true_torch, y_pred_torch, average='macro')
print('numpy unique')
%timeit np.unique(y_true)
print('torch unique')
%timeit torch.unique(y_true_torch)
print('numpy')
%prun -s cumulative jaccard_score(y_true, y_pred, average='macro')
print('torch')
with sklearn.config_context(array_api_dispatch=True):
%prun -s cumulative jaccard_score(y_true_torch, y_pred_torch, average='macro')Output: |
|
@lesteve Thanks for the benchmarks. |
|
Maybe another part of the performance difference: the magic of computing unique values once and saving the result in the dtype metadata is numpy-specific. If Having said that I guess computing metrics is unlikely to be the bottleneck so I would say this is low priority for now. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Reference Issues/PRs
Towards #26024
What does this implement/fix? Explain your changes.
Any other comments?