Possibility of optimisation in Focal Loss implementation

[vedantdalimkar]

I was training a multiclass model using focal loss provided by SMP. However, it was proving to be a bottleneck in training as switching torch's CrossEntropyLoss with smp's FocalLoss increased training time by a factor of 6x.

I think this may be caused by the suboptimal multiclass mode implementation where the function iterates over each class and calls the focal_loss_with_logits function seperately. This operation can be vectorised by calling focal_loss_with_logits once for all classes where target would be the one-hot encoding of the segmentation mask.

[vedantdalimkar]

Hi @qubvel. I have made a PR for this issue, please check and let me know if it needs any improvements! Thanks!