✨ Add EoMT Model || 🚨 Fix Mask2Former loss calculation

[yaswanth19]

What does this PR do?

Fixes #37171 and continuation of #37392

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

[Rocketknight1]

Image segmentation model so cc @qubvel @NielsRogge!

[yaswanth19]

@qubvel A rough draft is ready for inference 🤗 . Now i am adding support for training and they use mask_annealing to determine probability for attn masks. Is it required in this HF implementation also (I don't see for any other model) or are we fine with having a fixed prob for attn mask

[qubvel]

Hey @yaswanth19, do you mean the mask probability changes during the model training? Would be nice to have it but not sure it can be easily implemented tbh. Maybe we can just add a callback for a Trainer to change it? similar to learning rate callback (I mean not adding it to the Transformers actually, but to the docs/fine-tuning guide)

[yaswanth19]

do you mean the mask probability changes during the model training?

Yup exactly, and adding a trainer callback seems to be a good idea. I will check the feasibility of implementation and if simple enough then we can implement in the model itself else pivot to trainer callback.

[tommiekerssies]

@yaswanth19 Thanks for your great work!

A few thoughts:
• Mask annealing: Setting it to a fixed probability of 1 requires masked attention during inference, which we want to avoid. A fixed probability of 0 removes masked attention but does hurt performance. If mask annealing isn’t feasible right now, I’d suggest setting it to 0 for now. If you do implement it, note that intermediate mask predictions will be required, so the _predict method might need to move into EoMTEncoder. Just make sure intermediate masks are skipped at inference to avoid slowing things down.
• Testing: It might be good to add an integration test that verifies both the mask and class outputs for semantic, instance, and panoptic segmentation.
• Weight initialization: The current approach likely misinitializes parameters like query embeddings (e.g. std=0.02). All parameters outside the ViT should follow PyTorch defaults (e.g. nn.Embedding uses std=1.0).
• ViT backbone: Would it be simpler to use timm for the ViT? This allows loading pre-trained weights when training EoMT from scratch, avoids unnecessary code, and avoids unsupported use cases like training from a randomly initialized ViT.
• Loss function: Could we reuse the Mask2Former loss? Unless it conflicts with Transformers guidelines, this might reduce redundancy.

Let me know your thoughts.

[yaswanth19]

Thanks @tommiekerssies for your initial thoughts.

Mask annealing: Setting it to a fixed probability of 1 requires masked attention during inference, which we want to avoid.

• I am not sure abut mask annealing implementation compatibility with transformers natively. As I have said above, in the worst case we can set to 0 or use tariner callback if that's feasible.

Testing: It might be good to add an integration test

Yup, will add the complete test suite once I have a implementation ready.

Weight initialization: The current approach likely misinitializes parameters like query embeddings

Thanks for bringing this to my attention, I can make some correction to initialization later on when we have the end-to-end code ready. IMO, most of the user don't init from scratch and will either finetune it or just perform inference. But having said that I will look at timm implementation and will init in the same way

ViT backbone: Would it be simpler to use timm for the ViT?

Ideally yes 😅 But that's not the library coding standard (Don't want to introduce a hard dependency on timm). Also using timm backbone directly will not be compatible with all other features that HF ecosystem provides IMO. I am actually referring the HF VIT and timm implementation to get the best of both worlds and as to not introduce any bug.

Loss function: Could we reuse the Mask2Former loss?

Transformers has one model one file philosophy and because of that I have copied the Mask2Former loss completely here. It can be subjective call with Modular file in the sense we can expose Mask2Former loss and import it here for EoMT (Will require additional changes in mask2former) but that can be discussed during reviews with the core maintainer.

[tommiekerssies]

Thanks @tommiekerssies for your initial thoughts.

Mask annealing: Setting it to a fixed probability of 1 requires masked attention during inference, which we want to avoid.

• I am not sure abut mask annealing implementation compatibility with transformers natively. As I have said above, in the worst case we can set to 0 or use tariner callback if that's feasible.

Testing: It might be good to add an integration test

Yup, will add the complete test suite once I have a implementation ready.

Weight initialization: The current approach likely misinitializes parameters like query embeddings

Thanks for bringing this to my attention, I can make some correction to initialization later on when we have the end-to-end code ready. IMO, most of the user don't init from scratch and will either finetune it or just perform inference. But having said that I will look at timm implementation and will init in the same way

ViT backbone: Would it be simpler to use timm for the ViT?

Ideally yes 😅 But that's not the library coding standard (Don't want to introduce a hard dependency on timm). Also using timm backbone directly will not be compatible with all other features that HF ecosystem provides IMO. I am actually referring the HF VIT and timm implementation to get the best of both worlds and as to not introduce any bug.

Loss function: Could we reuse the Mask2Former loss?

Transformers has one model one file philosophy and because of that I have copied the Mask2Former loss completely here.

Thanks for the clarifications!

Regarding mask annealing, I agree that 0 for now is fine. That means effectively disabling masked attention and mask annealing, which is what the current code already does, so no changes needed on that front.

For weight initialization and the ViT backbone, I understand the constraints around using timm. In that case, I’d just make sure that the non-ViT parameters (query embeddings, mask MLP, upscale blocks, class head) aren’t using any custom initializations and instead follow PyTorch defaults. Should be a quick fix.

Let me know if you’d like me to look at any part in more detail.

[yaswanth19]

Hi @qubvel ,I’m working on refactoring the training logic for EoMT , and I’m running into a design challenge:

In the original single‐class implementation, they call _predict (which uses the class predictor head) on intermediate layer outputs to build attention masks. Because everything lives in one class, this is straightforward.

Refer: https://github.com/tue-mps/eomt/blob/c311b377d3189c976163e4ceb2156d90bb7db88f/models/eomt.py#L130

In our modular HF version, the encoder (EoMTEncoder) only runs the transformer blocks, and _predict (with mask_head, upscale_block, and class_predictor) lives in EoMTModel or EoMTForUniversalSegmentation. That separation means the encoder loop can’t access _predict , so we can’t reconstruct the original training flow.

I have two solutions in mind, LMK your thoughts on the below approaches and suggest any other better alternative:

1.) Club all the classes from EoMTEncoder into EomtForUniversalSegmentation, in this ways we can do all the processing in a single forward class.

2.) Move _predict which includes mask_head, upscale_block into EoMTEncoder and somehow pass the class_head. Pass these modules into the encoder class so that inside its forward loop it can call _predict, build the attention mask, and feed it into the next block. IMO this is a bit dirty and flow is tangled 😅 .

Here the _predict func is the same code which is used in Mask2Former for get mask_logits and class_logits from model output.

[tommiekerssies]

Yep, The model keys conversion is done offline, using the convert_eomt_to_hf.py script. It's not on the fly mapping , so we need to convert the checkpoints and store them separately.

Cmd which I am using: python src/transformers/models/eomt/convert_eomt_to_hf.py --hf_repo_id "" --output_dir "" --safe_serialization

Yep, The model keys conversion is done offline, using the convert_eomt_to_hf.py script. It's not on the fly mapping , so we need to convert the checkpoints and store them separately.

Cmd which I am using: python src/transformers/models/eomt/convert_eomt_to_hf.py --hf_repo_id "" --output_dir "" --safe_serialization

Understood. Is there an easy way to make it online, or is this not the way it should be for Transformers library?

[yaswanth19]

AFAIK you can directly write to hub using output_hub_path arg in conversion script. In that case only update would be readme afterwards depending on the checkpoint

[tommiekerssies]

AFAIK you can directly write to hub using output_hub_path arg in conversion script. In that case only update would be readme afterwards depending on the checkpoint

What I meant is that the from_pretrained() will take the checkpoint and convert it every time you call it. So we can keep the current checkpoints without changing anything. But maybe this is not the way it should be in the library.

[yaswanth19]

What I meant is that the from_pretrained() will take the checkpoint and convert it every time you call it. So we can keep the current checkpoints without changing anything. But maybe this is not the way it should be in the library.

Ahh, AFAIK I don't think that happens. We have to save the HF compatible ckpts in hub manually. CC: @yonigozlan here if there is a better way or can the conversion happen implicitly 🤔 .

[yonigozlan]

Hey @yaswanth19 and @tommiekerssies !

What I meant is that the from_pretrained() will take the checkpoint and convert it every time you call it. So we can keep the current checkpoints without changing anything. But maybe this is not the way it should be in the library.

No that's not how the conversion work in Transformers, the conversion script is only executed manually, and not called by from_pretrained().

I think the best thing to do here, since the current checkpoint doesn't use model.safetensors, is simply to upload all the files from @yaswanth19 checkpoints to the current tue-mps checkpoints, as there shouldn't be any conflicting files (we should be able to merge the two config.json without issues from what I'm seeing).

That means that when instantiating the model with from_pretrained() without using the trust_remote_code argument, you will get the Transformers version of the model (pytorch_model.bin won't be downloaded) and you'll get the current pytorch_model.bin weights when using it (not sure if trust_remote_code is used in the original repo currently @tommiekerssies ? Or if the checkpoints are just used to manually download the weights). In any case, uploading the Transformers checkpoint to the current checkpoints should work and that way we avoid confusion by having distinct checkpoints.

[yaswanth19]

Thank you @yonigozlan , Indeed we can keep both in the same repo. I was of the opinion that the bin file would be loaded if both are in same repo 😅 or error will pop up but hah model.safetensors would be loaded 👌.

[tommiekerssies]

Sounds like a good solution! @yaswanth19 I’ve given you write access to the org. Current configs aren’t used in our code at all, so they can be overwritten with whatever is needed for the Transformers library. So long as the pytorch_model.bin stay there the original code should keep working.

[yaswanth19]

Thanks @tommiekerssies , I have updated those 4 checkpoints which I had. Can you please check everything at your end - that you are able to load the same checkpoint for your implementation. You might need to add trust_remote_code=True or only download model.bin and load it 🤔 . Also made the suggested change in model's forward - please have a look.

[tommiekerssies]

Thanks @tommiekerssies , I have updated those 4 checkpoints which I had. Can you please check everything at your end - that you are able to load the same checkpoint for your implementation. You might need to add trust_remote_code=True or only download model.bin and load it 🤔 . Also made the suggested change in model's forward - please have a look.

Sorry, I made one final small comment about that change in model's forward ;)
Regarding the checkpoints, I see the model.bin are still there, so that works for our code.
For the remaining checkpoints, would you like me to convert them too using the conversion script in this PR and upload their respective files (safetensors, readme and configs with the right settings)?
Maybe it's a good idea to add a test case for each, so we know they all work? I guess just checking if they can be loaded is enough.

[yaswanth19]

I made one final small comment about that change in model's forward ;)

Yep, Done @tommiekerssies 🎉

Maybe it's a good idea to add a test case for each, so we know they all work? I guess just checking if they can be loaded is enough.

That should be covered in integration tests. I have tested the inference by pulling tue-mps updated repo and it's working fine 🤗 .

would you like me to convert them too using the conversion script in this PR

yes please, that would be great 😅 . I haven't tried the latest 7B one but assuming it's just a scaled up arc and in that case the conversion script should work fine

[yaswanth19]

@Cyrilvallez Please do the honours and merge the PR 🚀 🤗 .

[Cyrilvallez]

Yes, LGTM! Thanks again for the clean PR! 🤗

[qubvel]

@yaswanth19, congrats on the model merged 🎉 🎉 🎉

[yaswanth19]

Thanks @qubvel 🙏 🤗

[yonigozlan]

Congrats 🎉!