Unable to use torch scripting to export Mask2Former model
Author: rayryeng
Pull request: no
Labels: name: Good Difficult Issue

### System Info

- `transformers` version: 4.34.1
- Platform: Linux-5.4.0-100-generic-x86_64-with-glibc2.17
- Python version: 3.8.18
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes, a GTX 1080 with 8 GB of VRAM
- Using distributed or parallel set-up in script?: No.

### Who can help?

@amyeroberts 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I am attempting to export the [Mask2Former model available in huggingface](https://huggingface.co/facebook/mask2former-swin-base-coco-panoptic) through [`torch.jit.script`](https://pytorch.org/docs/stable/generated/torch.jit.script.html). Here's a minimal reproducible example:

```python
import torch
from transformers import Mask2FormerForUniversalSegmentation
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Mask2FormerForUniversalSegmentation.from_pretrained(
    "facebook/mask2former-swin-base-coco-panoptic", torchscript=True
    ).to(device)

scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'mask2former.pt')
```

By doing this, I get the following error using torch scripting (path to the offending file has been obfuscated for brevity):

```
torch.jit.frontend.NotSupportedError: Comprehension ifs are not supported yet:
  File "/home/.../huggingface/lib/python3.8/site-packages/transformers/models/mask2former/modeling_mask2former.py", line 2559
        if not return_dict:
            output = tuple(v for v in output.values() if v is not None)
            if loss is not None:
                output = ((loss)) + output
```

As a hack, I've changed my local installation so that comprehension ifs are removed:

```
    if not return_dict:
        outputs = []
        for v in output.values():
            if v is not None:
                outputs.append(v)
        output = tuple(outputs)    
```

This also occurs at line 2306 in the same file, so I've made the same changes there. Once I fix this, there is an error in the forward method for the SWIN backbone:

```
RuntimeError: 
'Optional[Tensor]' object has no attribute or method 'shape'.:
  File "/home/.../anaconda3/envs/huggingface/lib/python3.8/site-packages/transformers/models/swin/modeling_swin.py", line 313
    def forward(self, pixel_values: Optional[torch.FloatTensor]) -> Tuple[torch.Tensor, Tuple[int]]:
        _, num_channels, height, width = pixel_values.shape
                                         ~~~~~~~~~~~~~~~~~~ <--- HERE
        if num_channels != self.num_channels:
            raise ValueError(
```

The forward method for the SWIN backbone is confusing, as the input type is declared to be `Optional` but the output type is not. The definition of this method clearly indicates that a concrete tuple is to be returned.

As a final experiment, I've removed the `Optional` type declaration and tried to export it one more time:

```
aten::pad(Tensor self, SymInt[] pad, str mode="constant", float? value=None) -> Tensor:
Expected a value of type 'List[int]' for argument 'pad' but instead found type 'Tuple[int, Tensor]'.
:
  File "/home/.../anaconda3/envs/huggingface/lib/python3.8/site-packages/transformers/models/swin/modeling_swin.py", line 306
        if width % self.patch_size[1] != 0:
            pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
            pixel_values = nn.functional.pad(pixel_values, pad_values)
                           ~~~~~~~~~~~~~~~~~ <--- HERE
        if height % self.patch_size[0] != 0:
            pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
'SwinPatchEmbeddings.maybe_pad' is being compiled since it was called from 'SwinPatchEmbeddings.forward'
  File "/home/.../anaconda3/envs/huggingface/lib/python3.8/site-packages/transformers/models/swin/modeling_swin.py", line 319
            )
        # pad the input to be divisible by self.patch_size, if needed
        pixel_values = self.maybe_pad(pixel_values, height, width)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        embeddings = self.projection(pixel_values)
        _, _, height, width = embeddings.shape
```

It seems that what is being put into the forward pass is not, in fact, a `torch.Tensor` when being scripted.

Is torch scripting this model not supported at this time or am I missing something?

### Expected behavior

The model successfully being exported to disk with torch scripting.

Comments:
Hi @rayryeng, thanks for raising this issue and for the detailed PR description, it really helps us address issues 🤗 

No, the model isn't compatible with `torch.script` yet unfortunately. I think the correct type is being passed, it's just the model has several assumptions / logic flows / incorrect typing which are OK for eager execution but are not when trying to compile the functions. 

I don't have the bandwidth at the moment to convert the modeling files to make them compatible, but very happy to review any PRs from anyone in the community who would like to tackle this issue.  

Hi @amyeroberts - Thanks so much! My original intent in using torch scripting was because I am working on a real-time system, and it's a bit prohibitive for us to use a traced model as it will inevitably require some warmup for the model to compile before it can be used. I figured that scripting would eliminate the warmup required. For now, I think I can develop a workaround where we can warm up the model in a separate thread while we launch other things. Thankfully, panoptic segmentation is not the first thing to be done when we run our application. This is not a blocker for us, but more of a nice to have now.

If there's currently no bandwidth to do this, I'd be happy to try a first attempt, as it would benefit me as well as the community. I appreciate you having a look at this either way!

@rayryeng Great! 🤗 Having this enabled for this model will be very impactful for lots of users in the community. Feel free to open a PR and ping me for review or any other questions in the meantime 

Having the model exported in ONNX/TorchScript would be indeed awesome!