🚨🚨🚨 fix(Mask2Former): torch export 🚨🚨🚨
Author: philkuz
Pull request: yes
Labels: name: Vision

name: run-slow

name: ExecuTorch

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #34390 (issue)

Mask2Former modeling had a set of issues that prevented `torch.export` from working. This PR addresses ~3 individual problems that prevented the model from working. I took direction from a few PRs
that made similar changes:
1. [fix(Wav2Vec2ForCTC): torch export #34023](https://github.com/huggingface/transformers/pull/34023) for the attention mask modification to prevent a RuntimeError during torch.export
2. [Optim deformable detr #33600](https://github.com/huggingface/transformers/pull/33600) for passing around [spatial_shapes as a list](https://github.com/huggingface/transformers/pull/33600/files#diff-137bda311c71a6ad1765c311a04366dc1b547213c95795752ddb11e39466bbabR701) and as a tensor [and for the if statement fix](https://github.com/huggingface/transformers/pull/33600/files#diff-137bda311c71a6ad1765c311a04366dc1b547213c95795752ddb11e39466bbabR711)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
@amyeroberts, @qubvel, @ylacombe (tagging because you looked at #34023)

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->
 ## Testing:
 Run this once in `main` and then once with this branch. Ensure the testing.all_close() works.
 ```python
 import os

import numpy as np
import requests
import torch
from PIL import Image
from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation


def get_test_image():
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    image = Image.open(requests.get(url, stream=True).raw)
    return image


processor = AutoImageProcessor.from_pretrained(
    "facebook/mask2former-swin-base-coco-panoptic"
)
model = Mask2FormerForUniversalSegmentation.from_pretrained(
    "facebook/mask2former-swin-base-coco-panoptic"
)


results = model(**processor(images=get_test_image(), return_tensors="pt"))

results = processor.post_process_panoptic_segmentation(results)

results_file = "segmentation_results.pt"
if os.path.exists(results_file):
    old_segmentation = torch.load(results_file)
    np.testing.assert_array_almost_equal(results[0]["segmentation"], old_segmentation)
else:
    torch.save(results[0]["segmentation"], results_file)

scripted_model = torch.export.export(model, args=(torch.randn(1, 3, 800, 1280),))

```

I also visually compared the output masks and they look the same

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34393). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hey @yonigozlan if you have a chance, I would love another review of this PR after addressing your comments!
Thanks

Looks great to me thanks again for fixing this! I think we should add a short comment just above the definition of `level_start_index_list` explaining why we do this (iterating over a tensor breaks torch.compile/export).
Please rebase on main and then ask a core maintainer for a final review :).

Thanks Yoni!

Pinging @amyeroberts and @qubvel for a core maintainer review!

> Thanks for working on this!
> 
> It would be also nice to have a test that export works, maybe smth similar to [this](https://github.com/huggingface/transformers/commit/5392f12e1614383270ae8df524415a1f6b555773).

Added an export test and addressed your other comments. Please let me know if you would like me to modify anything else!

> Thanks for a quick iteration! A few more comments.
> 
> Please, push empty commit with message`[run_slow] mask2former` at the end to trigger all slow tests run for mask2former model

You got it! I think I need maintainer approval to run the slow workflows

soft ping @ArthurZucker, regarding your last comment, I'm not sure it can be done without breaking changes, but as soon as it is the internal module of the model I suppose it is fine to change its signature. 🚨 can be added to PR message if it is still required

> soft ping @ArthurZucker, regarding your last comment, I'm not sure it can be done without breaking changes, but as soon as it is the internal module of the model I suppose it is fine to change its signature. 🚨 can be added to PR message if it is still required

Hi @ArthurZucker @qubvel, 

Hope you're doing well, could you find some time today or tomorrow to provide guidance on this PR? Would love to check it off of my list!

> > soft ping @ArthurZucker, regarding your last comment, I'm not sure it can be done without breaking changes, but as soon as it is the internal module of the model I suppose it is fine to change its signature. 🚨 can be added to PR message if it is still required
> 
> Hi @ArthurZucker @qubvel,
> 
> Hope you're doing well, could you find some time today or tomorrow to provide guidance on this PR? Would love to check it off of my list!

Hi @ArthurZucker @qubvel any chance we can move forward with this PR this week? Happy to do whatever you would like, just need to get guidance on what you would like to do here.

Hi @philkuz, sorry for the delay, the team was on the offsite this week. I will ping Arthur to get it reviewed and merged. Thanks for the patience 

> Hi @philkuz, sorry for the delay, the team was on the offsite this week. I will ping Arthur to get it reviewed and merged. Thanks for the patience

Thank you, Pavel!

cc @guangy10 

@philkuz Awesome! Thanks for expanding the export coverage to more models! 🚀 🚀 🚀