Focal loss optimisation#1236
Open
vedantdalimkar wants to merge 4 commits intoqubvel-org:mainfrom
Open
Conversation
Contributor
Author
|
@qubvel Gentle reminder, if you can take a look at this it would be great! |
qubvel
reviewed
Feb 19, 2026
Collaborator
qubvel
left a comment
There was a problem hiding this comment.
Sorry for the delay, can you please make sure formatting is passing to run the tests. Also it would be nice to add some test case to make sure it works as expected, thanks!
Comment on lines
+77
to
+79
| y_true[y_true == self.ignore_index] = num_classes | ||
| y_true_one_hot = torch.nn.functional.one_hot(y_true,num_classes = num_classes + 1) | ||
| y_true_one_hot = y_true_one_hot[ ... , : -1] |
Collaborator
There was a problem hiding this comment.
Suggested change
| y_true[y_true == self.ignore_index] = num_classes | |
| y_true_one_hot = torch.nn.functional.one_hot(y_true,num_classes = num_classes + 1) | |
| y_true_one_hot = y_true_one_hot[ ... , : -1] | |
| y_true[y_true == self.ignore_index] = num_classes | |
| y_true_one_hot = torch.nn.functional.one_hot(y_true, num_classes = num_classes + 1) | |
| y_true_one_hot = y_true_one_hot[ ... , :-1] |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
This PR addresses #1235
The current focal loss implementation iterates over each class and calculates focal loss in a class-wise manner. This is slightly inefficient and can be optimised by vectorising the loss computation in multiclass mode. Also, the current implementation uses expensive masking operations for filtering out pixels belonging to
ignore_indexclassI have also attached a notebook that benchmarks the new approach against the old one. The time improvement is significant, often speeding up the code by more than 2x! The notebook also shows that the output of the new function is consistent with the new one.
@qubvel let me know if I need to add anymore tests.