Dice loss module causes `TypeError: forward() got an unexpected keyword argument 'ignore_index'` in v1.1.0

[ff98li]

Checklist

1. I have searched related issues but cannot get the expected help. ☑️

• Related issue TypeError: forward() got an unexpected keyword argument 'ignore_index' #201

2. The bug has not been fixed in the latest version. ☑️

Describe the bug
Attempted to train SegFormer using Dice loss + CE loss; the following error occurred after updating to mmsegmentation v1.1.0:

Traceback (most recent call last):
  File "/home/fei/pkgs/mmsegmentation/tools/train.py", line 104, in <module>
    main()
  File "/home/fei/pkgs/mmsegmentation/tools/train.py", line 100, in main
    runner.train()
  File "/home/fei/software/mambaforge/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/runner.py", line 1701, in train
    model = self.train_loop.run()  # type: ignore
  File "/home/fei/software/mambaforge/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/loops.py", line 278, in run
    self.run_iter(data_batch)
  File "/home/fei/software/mambaforge/envs/openmmlab/lib/python3.8/site-packages/mmengine/runner/loops.py", line 301, in run_iter
    outputs = self.runner.model.train_step(
  File "/home/fei/software/mambaforge/envs/openmmlab/lib/python3.8/site-packages/mmengine/model/base_model/base_model.py", line 114, in train_step
    losses = self._run_forward(data, mode='loss')  # type: ignore
  File "/home/fei/software/mambaforge/envs/openmmlab/lib/python3.8/site-packages/mmengine/model/base_model/base_model.py", line 326, in _run_forward
    results = self(**data, mode=mode)
  File "/home/fei/software/mambaforge/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/fei/pkgs/mmsegmentation/mmseg/models/segmentors/base.py", line 94, in forward
    return self.loss(inputs, data_samples)
  File "/home/fei/pkgs/mmsegmentation/mmseg/models/segmentors/encoder_decoder.py", line 176, in loss
    loss_decode = self._decode_head_forward_train(x, data_samples)
  File "/home/fei/pkgs/mmsegmentation/mmseg/models/segmentors/encoder_decoder.py", line 137, in _decode_head_forward_train
    loss_decode = self.decode_head.loss(inputs, data_samples,
  File "/home/fei/pkgs/mmsegmentation/mmseg/models/decode_heads/decode_head.py", line 262, in loss
    losses = self.loss_by_feat(seg_logits, batch_data_samples)
  File "/home/fei/pkgs/mmsegmentation/mmseg/models/decode_heads/decode_head.py", line 324, in loss_by_feat
    loss[loss_decode.loss_name] = loss_decode(
  File "/home/fei/software/mambaforge/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
TypeError: forward() got an unexpected keyword argument 'ignore_index'

Reproduction

1. What command or script did you run?

Added Dice loss to my config for SegFormer as below:

  model = dict(
   ## Don't use data_preprocessor in predefined config
   data_preprocessor = data_preprocessor,
   pretrained = None,
   decode_head = dict(
       num_classes = 13, 
       loss_decode = [
           dict(
               type = "CrossEntropyLoss",
               loss_name = "loss_ce",
               loss_weight = 2.0
           ),
           dict(
               type = 'DiceLoss',
               loss_name = 'loss_dice',
               loss_weight = 1.0
           )
       ]
   )
)

I believe this issue is directly related to the Dice loss as after removing the Dice loss the error no longer occurs.

2. Did you make any modifications on the code or config? Did you understand what you have modified?

Yes, as described above. It used to work with dev-1.* before I updated mmseg to the main branch version.

3. What dataset did you use?
   🤐 Custom datasets, but used to work with dev-1.*
   Environment

sys.platform: linux
Python: 3.8.16 | packaged by conda-forge | (default, Feb  1 2023, 16:01:55) [GCC 11.3.0]
CUDA available: True
numpy_random_seed: 2147483648
GPU 0,1: NVIDIA RTX A6000
CUDA_HOME: /usr/local/cuda
NVCC: Cuda compilation tools, release 11.6, V11.6.124
GCC: gcc (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
PyTorch: 2.0.0
PyTorch compiling details: PyTorch built with:
  - GCC 9.3
  - C++ Version: 201703
  - Intel(R) oneAPI Math Kernel Library Version 2022.1-Product Build 20220311 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.7.3 (Git Hash 6dbeffbae1f23cbbeae17adb7b5b13f1f37c080e)
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - LAPACK is enabled (usually provided by MKL)
  - NNPACK is enabled
  - CPU capability usage: AVX2
  - CUDA Runtime 11.7
  - 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.5
  - Magma 2.6.1
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.7, CUDNN_VERSION=8.5.0, CXX_COMPILER=/opt/rh/devtoolset-9/root/usr/bin/c++, CXX_FLAGS= -D_GLIBCXX_USE_CXX11_ABI=0 -fabi-version=11 -Wno-deprecated -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wunused-local-typedefs -Wno-unused-parameter -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Werror=cast-function-type -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_DISABLE_GPU_ASSERTS=ON, TORCH_VERSION=2.0.0, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF, 

TorchVision: 0.15.0
OpenCV: 4.7.0
MMEngine: 0.7.0
MMSegmentation: 1.1.0+0079076

[callmewenhao]

i meet the same error🤣

[yeedrag]

I had similar problems when trying to use multiple losses too

[longphamkhac]

Beside the ignore_index, I think the dice loss implementation has it own logical problem as well. The output_preds has shape [bs, num_classes, H, W] but the target's shape is [bs, H, W], which is cannot be the same shape when flattenning. I think the target has to pass through the one hot encoding first.

[okdha1234]

facing similar issue. Any solution to this ,?
@yeedrag @callmewenhao @longphamkhac @kahkeng

[stone-cloud]

I'm also having a similar issue when using STDC-v2

[yeedrag]

facing similar issue. Any solution to this ,? @yeedrag @callmewenhao @longphamkhac @kahkeng

We solved it by just writing a new loss that calculates and combines both losses together, ugly but it works.

[DXY121]

me too.....

[xiexinch]

facing similar issue. Any solution to this ,? @yeedrag @callmewenhao @longphamkhac @kahkeng

We solved it by just writing a new loss that calculates and combines both losses together, ugly but it works.

Hi @yeedrag,
Could you create a PR to fix this problem? The dice loss was modified by this PR #3002, the implementation is the same as mmdet.

[yeedrag]

facing similar issue. Any solution to this ,? @yeedrag @callmewenhao @longphamkhac @kahkeng

We solved it by just writing a new loss that calculates and combines both losses together, ugly but it works.

Hi @yeedrag, Could you create a PR to fix this problem? The dice loss was modified by this PR #3002, the implementation is the same as mmdet.

Hi @xiexinch, our solution was to just make a new loss called focal_dice_loss, where we calculate the two losses inside. We didn't modify it in any way.
I was interested in fixing the problem though, so I did a few changes and now the dice loss can run properly:
I added ignore_index inside the forward function (although not used), and I added one hot encoding for pred to ensure it has the same dim as target ( which is [N, num_classes, H, W as @longphamkhac mentioned).
The problem is I'm not familiar with contributing to open sources, and this is basically my first time. How can I check my code is absolutely correct (I ran pytest on test_dice_loss.py and it passed) and what procedures do I need to do to make the pr?
I have also noticed a few weird things. The first one is ignore_index in dice loss isn't actually used although it should. also, the default value of ignore_index across losses is inconsistent with a few as -100 and some as 255.
Thanks in advance!

[xiexinch]

Hi @yeedrag,
Thanks for your interest, just follow this workflow, create a pull request, then we can review and merge your code.

[sasniper007]

I modified the dice_loss.py like below

def forward(self, pred, target, weight=None, avg_factor=None, reduction_override=None, ignore_index=255, **kwargs):
it works well

[yeedrag]

Hi @yeedrag, Thanks for your interest, just follow this workflow, create a pull request, then we can review and merge your code.

Hey @xiexinch, quick question: Im trying to commit my changes, and the pre-commit update-model-index hook modifies all the metafile.yaml files so I couldn't pass the pre-commit check if I only committed the dice_loss.py. Should I also commit the yaml files where the only change is the file paths are modified to be local, or is there a workaround? Sorry I'm not familiar with this and Thanks in advance.
(I'm not sure whether this should be opened as a new issue, as I think this is a rather insignificant problem)

[ff98li]

@yeedrag Thank you so much for the fix! I'm closing this issue now as the fix has been merged.