Repository: open-mmlab/mmpretrain

Issue:
Focal loss for single label classification?
### Checklist
- I have searched related issues but cannot get the expected help.
- I have read related documents and don't know what to do.

### Describe the question you meet

I'm trying to train a resnet18 with a LinearClsHead using the focal loss for a single label binary classification problem, but the current focal loss implementation raises an assertion error as the predictions have shape [batch_size, 2] and target has shape [batch_size]. According to the documentation the implemented focal loss only works for multilabel tasks. Is there any way of using it also for single label tasks? 

### Post related information
1. The output of `pip list | grep "mmcv\|mmcls\|^torch"`
mmcls                     0.17.0
mmcv-full                 1.3.9
torch                     1.9.0
torch-model-archiver      0.4.1
torchvision               0.10.0

2. This is the model config I'm using:
```python
# type: ignore
model = dict(
    type="ImageClassifier",
    backbone=dict(
        type="ResNet", depth=18, num_stages=4, out_indices=(3,), style="pytorch"
    ),
    neck=dict(type="GlobalAveragePooling"),
    head=dict(
        type="LinearClsHead",
        num_classes=2,
        in_channels=512,
        loss=dict(type="FocalLoss", loss_weight=1.0),
        topk=(1),
    ),
)

load_from = "/media/VA/pretrained_weights/mmcls/resnet18_batch256_20200708-34ab8f90.pth"
```
3. I am getting the following error during training:
```
File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcls/apis/train.py", line 164, in train_model
    runner.run(data_loaders, cfg.workflow)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcv/runner/epoch_based_runner.py", line 127, in run
    epoch_runner(data_loaders[i], **kwargs)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcv/runner/epoch_based_runner.py", line 50, in train
    self.run_iter(data_batch, train_mode=True, **kwargs)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcv/runner/epoch_based_runner.py", line 29, in run_iter
    outputs = self.model.train_step(data_batch, self.optimizer,
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcv/parallel/data_parallel.py", line 67, in train_step
    return self.module.train_step(*inputs[0], **kwargs[0])
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcls/models/classifiers/base.py", line 146, in train_step
    losses = self(**data)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcv/runner/fp16_utils.py", line 98, in new_func
    return old_func(*args, **kwargs)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcls/models/classifiers/base.py", line 90, in forward
    return self.forward_train(img, **kwargs)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcls/models/classifiers/image.py", line 110, in forward_train
    loss = self.head.forward_train(x, gt_label)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcls/models/heads/linear_head.py", line 53, in forward_train
    losses = self.loss(cls_score, gt_label)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcls/models/heads/cls_head.py", line 46, in loss
    loss = self.compute_loss(cls_score, gt_label, avg_factor=num_samples)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcls/models/losses/focal_loss.py", line 106, in forward
    loss_cls = self.loss_weight * sigmoid_focal_loss(
  File "/media/data/miniconda3/envs/ai-project-screen_classification/lib/python3.8/site-packages/mmcls/models/losses/focal_loss.py", line 36, in sigmoid_focal_loss
    assert pred.shape == \
AssertionError: pred and target should be in the same shape.
```
Thank you!


Discussion:
We will try to reproduce and fix it recently
Hello, by now, the `FocalLoss` is only for multi-label tasks, refers to https://github.com/open-mmlab/mmclassification/pull/131, and we will support single-label task in the future.
@juanjompz Hi. I met the same problem and I think I sloved it.
One-Hot is the key.
You can insert this code in line 37, focal_loss.py (before `assert pred.shape ==...`)
code:
```
    from .utils import convert_to_one_hot
    if target.dim() == 1 or (target.dim() == 2 and target.shape[1] == 1):
        target = convert_to_one_hot(target.view(-1, 1), pred.shape[-1])
```
Then, mmcls can work successfully.
Thank you @zhongqiu1245, indeed, converting target to one hot format fixes the problem.  Are you going to pull request these changes?

Resolving patch:
diff --git a/mmcls/models/losses/focal_loss.py b/mmcls/models/losses/focal_loss.py
--- a/mmcls/models/losses/focal_loss.py
+++ b/mmcls/models/losses/focal_loss.py
@@ -3,7 +3,7 @@
 import torch.nn.functional as F
 
 from ..builder import LOSSES
-from .utils import weight_reduce_loss
+from .utils import convert_to_one_hot, weight_reduce_loss
 
 
 def sigmoid_focal_loss(pred,
@@ -88,7 +88,7 @@ def forward(self,
         Args:
             pred (torch.Tensor): The prediction with shape (N, \*).
             target (torch.Tensor): The ground truth label of the prediction
-                with shape (N, \*).
+                with shape (N, \*), N or (N,1).
             weight (torch.Tensor, optional): Sample-wise loss weight with shape
                 (N, \*). Defaults to None.
             avg_factor (int, optional): Average factor that is used to average
@@ -103,6 +103,8 @@ def forward(self,
         assert reduction_override in (None, 'none', 'mean', 'sum')
         reduction = (
             reduction_override if reduction_override else self.reduction)
+        if target.dim() == 1 or (target.dim() == 2 and target.shape[1] == 1):
+            target = convert_to_one_hot(target.view(-1, 1), pred.shape[-1])
         loss_cls = self.loss_weight * sigmoid_focal_loss(
             pred,
             target,


Test patch:
diff --git a/tests/test_metrics/test_losses.py b/tests/test_metrics/test_losses.py
--- a/tests/test_metrics/test_losses.py
+++ b/tests/test_metrics/test_losses.py
@@ -161,6 +161,14 @@ def test_focal_loss():
     # test focal_loss with weight
     assert torch.allclose(
         loss(cls_score, label, weight=weight), torch.tensor(0.8522 / 2))
+    # test focal loss for single label task
+    cls_score = torch.Tensor([[5, -5, 0], [5, -5, 0]])
+    label = torch.Tensor([0, 1])
+    weight = torch.tensor([0.5, 0.5])
+    assert torch.allclose(loss(cls_score, label), torch.tensor(0.86664125))
+    # test focal_loss single label with weight
+    assert torch.allclose(
+        loss(cls_score, label, weight=weight), torch.tensor(0.86664125 / 2))
 
 
 def test_label_smooth_loss():