I want an example code for Focal loss in PyTorch for a model with three class prediction. My model outputs 3 probabilities.
Sentiment_LSTM( (embedding): Embedding(19612, 400) (lstm): LSTM(400, 512, num_layers=2, batch_first=True, dropout=0.5) (dropout): Dropout(p=0.5, inplace=False) (fc): Linear(in_features=512, out_features=3, bias=True) (sig): Sigmoid() )
My class distribution is highly imbalanced. So I want to try focal loss so that the minor class accuracy is improved.
I currently used loss function defined in https://www.kaggle.com/c/tgs-salt-identification-challenge/discussion/65938 But it didn't help.
The original paper(https://arxiv.org/abs/1708.02002) only consider binary classification. How do I extend it to the multi-class scenario?