# Focal loss for imbalanced multi class classification in Pytorch

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?