Add MCCLoss implementation for binary image segmentation #108
Conversation
|
@kakumarabhishek Hi, and thanks for your PR! |
pytorch_toolbelt/losses/mcc.py
Outdated
| batch_size = y_true.shape[0] | ||
|
|
||
| y_true = y_true.view(batch_size, 1, -1) | ||
| y_pred = y_pred.view(batch_size, 1, -1) |
There was a problem hiding this comment.
Here y_pred are logits (-inf,+inf) accroding to docstring, but underlying loss treat them as bounded (0,1) value (E.g .sigmoid() missing?)
There was a problem hiding this comment.
Thanks for this comment. You are right - in my usage, I always applied the loss after a sigmoid operation to the network output, but you are right, this should be an explicit choice. Updated the function to accommodate this choice.
pytorch_toolbelt/losses/mcc.py
Outdated
| tp = torch.sum(torch.mul(y_pred, y_true)) + self.eps | ||
| tn = torch.sum(torch.mul((1 - y_pred), (1 - y_true))) + self.eps | ||
| fp = torch.sum(torch.mul(y_pred, (1 - y_true))) + self.eps | ||
| fn = torch.sum(torch.mul((1 - y_pred), y_true)) + self.eps |
There was a problem hiding this comment.
Here tp,fn,fp,fn computed per the whole batch. Shouldn't we compute those per-sample? Given the shapes of y_true/y_pred it should be something like tp = torch.sum(..., dim=(1,2)) + eps).
This way, tp,fn,fp,fn will have shape of [B] and at very end you compute per-sample mean via loss.mean()
There was a problem hiding this comment.
Thank you. I updated the implementation to perform either a sample-level or batch-level reduction. This should address your comment.
Add option to calculate loss from either logits or predictions. Add option to perform sample-wise or batch-wise reduction. Add tests for `MCCLoss`.
|
Thanks @BloodAxe for reviewing the code. Both your comments are valid and I appreciate your inputs. Here are the changes:
Please let me know if you have any other concerns. |
|
@BloodAxe , can you please take a look at this? Thank you. |
|
@BloodAxe , any updates? |
2 participants
Added loss function for binary image segmentation with the Matthews Correlation Coefficient (MCC). Based on the reference implementation and the paper.