Repository: open-mmlab/mmdetection

Issue:
[Reimplementation] IndexError :argmax()  in refine_bboxes function of sparse-rcnn (latest version (3.x).)
### Prerequisite

- [X] I have searched [Issues](https://github.com/open-mmlab/mmdetection/issues) and [Discussions](https://github.com/open-mmlab/mmdetection/discussions) but cannot get the expected help.
- [X] I have read the [FAQ documentation](https://mmdetection.readthedocs.io/en/latest/faq.html) but cannot get the expected help.
- [X] The bug has not been fixed in the [latest version (master)](https://github.com/open-mmlab/mmdetection) or [latest version (3.x)](https://github.com/open-mmlab/mmdetection/tree/dev-3.x).

### 💬 Describe the reimplementation questions

When I try to train the sparse-rcnn series in the latest version (3. x) mmdet with my personal dataset (only 1 class), the following errors always occur:
```
  File "I:\mmdetection-3.x\mmdetection\mmdet\models\detectors\base.py", line 92, in forward
    return self.loss(inputs, data_samples)
  File "I:\mmdetection-3.x\mmdetection\mmdet\models\detectors\two_stage.py", line 187, in loss
    roi_losses = self.roi_head.loss(x, rpn_results_list,
  File "I:\mmdetection-3.x\mmdetection\mmdet\models\roi_heads\sparse_roi_head.py", line 344, in loss
    bbox_results = self.bbox_loss(
  File "I:\mmdetection-3.x\mmdetection\mmdet\models\roi_heads\sparse_roi_head.py", line 122, in bbox_loss
    bbox_results = self._bbox_forward(stage, x, rois, object_feats,
  File "I:\mmdetection-3.x\mmdetection\mmdet\models\roi_heads\sparse_roi_head.py", line 230, in _bbox_forward
    results_list = bbox_head.refine_bboxes(
  File "I:\mmdetection-3.x\mmdetection\mmdet\models\roi_heads\bbox_heads\bbox_head.py", line 645, in refine_bboxes
    cls_scores[:, :-1].argmax(1), labels)
IndexError: argmax(): Expected reduction dim 1 to have non-zero size.

Process finished with exit code 1
```
And the class_scores's shape is (100,1).100 is num_proposals. **I guess this 1 should be 2. Because there are two categories: foreground and background.**
![image](https://user-images.githubusercontent.com/49437933/201906282-e2090252-de70-44e9-a633-905910e9b8f8.png)

In the config file, I only modified the num_classes option in the head part and dataset address.

### Environment

```
sys.platform: win32
Python: 3.8.12 (default, Oct 12 2021, 03:01:40) [MSC v.1916 64 bit (AMD64)]
CUDA available: True
numpy_random_seed: 2147483648
GPU 0: NVIDIA GeForce RTX 3090
CUDA_HOME: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.4
NVCC: Cuda compilation tools, release 11.4, V11.4.152
GCC: n/a
PyTorch: 1.11.0
PyTorch compiling details: PyTorch built with:
  - C++ Version: 199711
  - MSVC 192829337
  - Intel(R) Math Kernel Library Version 2020.0.2 Product Build 20200624 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.5.2 (Git Hash a9302535553c73243c632ad3c4c80beec3d19a1e)
  - OpenMP 2019
  - LAPACK is enabled (usually provided by MKL)
  - CPU capability usage: AVX512
  - CUDA Runtime 11.3
  - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_61,code=sm_61;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_37,code=compute_37
  - CuDNN 8.2
  - Magma 2.5.4
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.3, CUDNN_VERSION=8.2.0, CXX_COMPILER=C:/cb/pytorch_1000000000000/work/tmp_bin/sccache-cl.exe, CXX_FLAGS=/DWIN32 /D_WINDOWS /GR /EHsc /w /bigobj -DUSE_PTHREADPOOL -openmp:experimental -IC:/cb/pytorch_1000000000000/work/mkl/include -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOCUPTI -DUSE_FBGEMM -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -DEDGE_PROFILER_USE_KINETO, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_VERSION=1.11.0, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=OFF, USE_MPI=OFF, USE_NCCL=OFF, USE_NNPACK=OFF, USE_OPENMP=ON, USE_ROCM=OFF,

TorchVision: 0.12.0
OpenCV: 4.1.2
MMEngine: 0.3.1
MMDetection: 3.0.0rc3+5b0d5b4
```

### Expected results

_No response_

### Additional information

1.In config: 
- changed the data and annotation path
- added metainfo to add my 'CLASSES'
- changed 'num_classes' from 80 to 1.

2.I also test mmdet3.0 with my dataset in cascade rcnn.And the second dim of `cls_scores` is 2. So I suppose is there some problems in `DIIHead`? Or just I use the wrong way to train sparse rcnn?


Discussion:
I think it might the sigmoid function of focal loss make `cls_scores[:, :-1]` just 1 element of dim 2 .So I change [Line 644](https://github.com/open-mmlab/mmdetection/blob/5b0d5b40d5c6cfda906db7464ca22cbd4396728a/mmdet/models/roi_heads/bbox_heads/bbox_head.py#L644)
from
```
        labels = torch.where(labels == self.num_classes,
                             cls_scores[:, :-1].argmax(1), labels)
```
to
```
        if self.num_classes == 1:
            labels = torch.where(labels == self.num_classes,
                                 0, labels)
        else:
            labels = torch.where(labels == self.num_classes,
                                 cls_scores[:, :-1].argmax(1), labels)
```
in `bbox_head.py`.
And it works.I'm not sure if this will have bad consequences (and if this way is right).
Can anyone help? QAQ
The DIIHead assumes the head use FocalLoss. This is a bit weird. Does your modification improve the performance?
@ZwwWayne ` cls_scores[:, :-1]` make the second dimension reduced by one (for other loss, it might means removing the `background` class.). So `cls_scores' shape: [100, 1] --> [100, 0]`.Then this error (` IndexError `) will appear. As shown in the following figure:
![image](https://user-images.githubusercontent.com/49437933/202971583-fcf2ff08-d3cf-4bb5-ab96-502bdb9ec506.png)
My modification can just make it work ( mAP50 70+, which is a normal value), but I don't think it's the fundamental way to solve this problem.
Thanks for your report! We found that this is a bug introduced in b21714cbc8bf815190c7b272c6644c9624dd590a
We will fix it ASAP.

Resolving patch:
diff --git a/mmdet/models/roi_heads/bbox_heads/bbox_head.py b/mmdet/models/roi_heads/bbox_heads/bbox_head.py
--- a/mmdet/models/roi_heads/bbox_heads/bbox_head.py
+++ b/mmdet/models/roi_heads/bbox_heads/bbox_head.py
@@ -640,9 +640,15 @@ def refine_bboxes(self, sampling_results: Union[List[SamplingResult],
             cls_scores = self.loss_cls.get_activation(cls_scores)
         if cls_scores.numel() == 0:
             return None
-
-        labels = torch.where(labels == self.num_classes,
-                             cls_scores[:, :-1].argmax(1), labels)
+        if cls_scores.shape[-1] == self.num_classes + 1:
+            # remove background class
+            cls_scores = cls_scores[:, :-1]
+        elif cls_scores.shape[-1] != self.num_classes:
+            raise ValueError('The last dim of `cls_scores` should equal to '
+                             '`num_classes` or `num_classes + 1`,'
+                             f'but got {cls_scores.shape[-1]}.')
+        labels = torch.where(labels == self.num_classes, cls_scores.argmax(1),
+                             labels)
 
         img_ids = rois[:, 0].long().unique(sorted=True)
         assert img_ids.numel() <= len(batch_img_metas)


Test patch: