search

MIC-DKFZ / nnUNet

Apache License 2.0
5.9k stars 1.76k forks source link

Adding weights for multi class segmentation #2008

Closed sgritter closed 6 months ago

sgritter commented 8 months ago

I am using the nnunetv2 for segmenting lung tissue and i have three different classes: aired tissue, non-aired tissue and background. I want to adjust the weights to make sure the model prioritizes to segment the aired tissue right. However I cannot find any way to do this. Is this possible?

GregorKoehler commented 6 months ago

Hi @sgritter , the cross entropy loss used in nnUNet (e.g. in this trainer https://github.com/MIC-DKFZ/nnUNet/blob/3808231601265aed98d70de4218fc845e185959d/nnunetv2/training/nnUNetTrainer/variants/loss/nnUNetTrainerCELoss.py#L8 ) does inherit from torch.nn.CrossEntropyLoss, which has a weight parameter. In the linked trainer, you see this argument set to None. But you can create your own trainer like this one, but instead set weights in the _build_loss method.

Let me know if this helped in your case!

GregorKoehler commented 6 months ago

Closing this issue for now, as it was stale for roughly a week. Please feel free to re-open if you're still facing this issue!