How to enable multi-heads outputs in segformer?
Author: JunMa11
Pull request: no

Dear all,

Thanks for the awesome work:)

I was wondering how to enable segformer to output two maps? E.g., one map for multi-class labels (default) and the other map for distance map

![1645718354(1)](https://user-images.githubusercontent.com/19947331/155560474-3769338e-b68d-4739-86ba-c26b6b8641c0.png)

https://github.com/huggingface/transformers/blob/35ecf99cc4e1046e5908f5933a4548e43a0f4f1d/src/transformers/models/segformer/modeling_segformer.py#L681

Comments:
Hi,

Thanks for your interest in SegFormer! Here's how you can easily create a SegFormer with custom loss:

```
import transformers import SegformerPreTrainedModel, SegformerForSemanticSegmentation
from transformers.modeling_outputs import SemanticSegmentationModelOutput
import torch.nn as nn

class CustomSegFormer(SegformerPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.segformer = SegformerForSemanticSegmentation.from_pretrained("...")

    def forward(self, pixel_values, labels=None):
        # forward to get raw logits
        outputs = self.segformer(pixel_values)
        logits = outputs.logits
        
        # next, do whatever you want with the logits to compute loss
        loss = None
        if labels is not None:
               # upsample logits here
               upsampled_logits = nn.functional.interpolate(logits, size=pixel_values.shape[-2:], mode="bilinear", align_corners=False)
               # calculate custom loss (or losses) here
               loss  = ...
 
        return SemanticSegmentationModelOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
```

A few notes here:
* I'm inheriting from `SegformerPreTrainedModel` rather than `nn.Module` since you can then also use nice utility methods such as `from_pretrained`, `save_pretrained` and `push_to_hub`. It will also take care of proper weights initialization as seen [here](https://github.com/huggingface/transformers/blob/7566734d6f6467f1442bf0494cb14108d72c61b3/src/transformers/models/segformer/modeling_segformer.py#L416).
* I'm returning a `SemanticSegmentationModelOutput`. This makes the model compatible with HuggingFace's Trainer class.

Hi @NielsRogge,

Thanks for your kind help very much.

The above demo is very useful to me but it only has one output branch.  I want the `segformer` output to have two branches and then, I can apply different loss functions to the two branches.

I'm a newcomer to hugging face, could you please give a quick check to my modification?

Thanks for your time very much.

```python
import transformers import SegformerPreTrainedModel, SegformerForSemanticSegmentation
from transformers.modeling_outputs import SemanticSegmentationModelOutput
import torch.nn as nn


class CustomSegFormer(SegformerPreTrainedModel):
    def __init__(self, config, in_channel, n_classes=3):
        super().__init__(config)
        self.segformer = SegformerForSemanticSegmentation.from_pretrained("...")
        self.logits_outputs = nn.Conv2d(in_channel, output_channel=n_classes, kernel_size=1, stride=1, padding=0)
        self.distance_outputs = nn.Conv2d(in_channel, output_channel=1, kernel_size=1, stride=1, padding=0)

    def forward(self, pixel_values, calss_labels, distance_labels):
        # forward to get raw logits
        outputs = self.segformer(pixel_values) # what's the number of output channels?
        logits_outputs = self.logits_outputs(outputs)
        distancemap_outputs = self.distance_outputs(outputs)        

        # next, do whatever you want with the logits to compute loss
        # upsample logits here; the predicted distance map is also needed to unsampled, right?
        upsampled_logits = nn.functional.interpolate(logits_outputs, size=pixel_values.shape[-2:], mode="bilinear", align_corners=False)
        upsampled_distancemaps = nn.functional.interpolate(distancemap_outputs , size=pixel_values.shape[-2:], mode="bilinear", align_corners=False)
        # calculate losses for logits and distance map respectively and sum the two losses
        loss_logits = nn.functional.cross_entropy(unsampled_logits, calss_labels)
        loss_dists  = torch.norm(distance_labels- upsampled_distancemaps)/torch.numel(distance_labels)
        loss = loss_logits + loss_dists

```


It depends on how you want to compute the logits and distance maps.

Note that the model (`SegformerForSemanticSegmentation`) returns a `SemanticSegmentationModelOutput`, which already contains `logits` as well as all intermediate `hidden_states` (feature maps). The latter will only be included in the outputs if you pass `output_hidden_states=True`;

```
outputs = self.segformer(pixel_values, output_hidden_states=True)
logits = outputs.logits
hidden_states = outputs.hidden_states
```
As you'll see, the `hidden_states` are a tuple of `torch.FloatTensor`, containing the intermediate feature maps.

How do you want the distance maps to be computed? Based on the features of the last layer of the model?