[SegFormer] Add support for segmentation masks with one label

[NielsRogge]

What does this PR do?

This PR makes it possible to fine-tune SegFormer in case you have a mask containing only a single value, i.e. your mask could look like [[255, 0], [0, 255]]. In this case, config.num_labels = 1 and the ignore index is 255.

If this works fine, then we can add it to any other model supported by AutoModelForSemanticSegmentation.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[dimkoss11]

How soon will this PR be done? And why using BCEWithLogitsLoss not Dice loss?

[jamilkarabaev]

when is this going to be closed. Our team would like to use it

[jamilkarabaev]

when this becomes active, how would one use it.

thank you nielsrogge for sending for review.

[jamilkarabaev]

hello @NielsRogge

hope I didnt disturb

I tried to peel off the classifier in pytorch and change the output channels to one, then manually compute the loss instead of getting it from the segformer huggingface object, so basically I just got the output and then did a dice loss myself.

So i tried to write bin seg myself kinda, and I started to get a bunch of negative values.

Any idea why that happened, and how I could fix it? I mean, I replaced the last conv2d layer.

[alaradirik]

Looks good to me! But I agree with the commenter about using Dice loss instead of BCEWithLogitsLoss.

[sgugger]

Looks good to me once the loss comment is resolved.

[aegonwolf]

Is this now automatically using dice loss when we set num_labels = 1? Maybe I missed it but it seems the documentation doesn't explain it.

[NielsRogge]

HHi @aegonwolf,

When config.num_labels = 1, the binary cross-entropy loss is used, as can be seen here:

transformers/src/transformers/models/segformer/modeling_segformer.py

Lines 813 to 817 in 1689aea

	elif self.config.num_labels == 1:
	valid_mask = ((labels >= 0) & (labels != self.config.semantic_loss_ignore_index)).float()
	loss_fct = BCEWithLogitsLoss(reduction="none")
	loss = loss_fct(upsampled_logits.squeeze(1), labels.float())
	loss = (loss * valid_mask).mean()

.

[nikolaJovisic]

What about a sigmoid function at the end or any way of scaling outputs to [0, 1] range?

[NielsRogge]

@nikolaJovisic the BCE loss uses sigmoid.

[hungphongtran-pixta]

@NielsRogge Thank you for your amazing work. However I am still struggling with this task, can you make an example notebook on how to finetune Segformer for 1 class only please? Thank you a lot!

[nikolaJovisic]

@NielsRogge modeling_segformer.py script is changed, but modeling_tf_segformer.py seems to be unchanged, I get an error there when I try to train model for 1 class only.

[pratikkorat26]

Any update on this task ? Any example notebook ?

[nikolaJovisic]

Any update on this task ? Any example notebook ?

Hi,

Back when I was working on this, the BinaryCrossentropy variant was merged and the issue was closed, but if I remember correctly it was done only for PyTorch version, so I implemented the Tensorflow version myself on a fork of this repository, you can find it on my profile, never made the effort to open a PR tho. Hope this might help. The fine-tuning itself was straightforward after that if I remember correctly.

Best wishes,
Nikola

[NielsRogge]

@nikolaJovisic would be great if you could open a PR for it!