Repository: huggingface/transformers
Pull request: Fix initialization of OneFormer

# 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 -->

Following #38061 and #38864, I noticed the same error for OneFormer due to duplicate initializations. I essentially applied the same fix as in #38864.

Code to reproduce the bug:
```python
from transformers import (
    AutoModelForImageClassification,
    OneFormerForUniversalSegmentation,
    OneFormerConfig,
)

backbone_name = "microsoft/resnet-18"
# backbone_name = "microsoft/swin-tiny-patch4-window7-224"

def params_match(params1, params2):
            return all((p1 == p2).all() for p1, p2 in zip(params1, params2))

# load pretrained backbone model
backbone_model = AutoModelForImageClassification.from_pretrained(backbone_name)

# load OneFormerConfig with a pretrained backbone
config = OneFormerConfig(
    backbone=backbone_name,
    use_pretrained_backbone=True,
)
oneformer = OneFormerForUniversalSegmentation(config)

# AssertionError: parameters don't match
assert params_match(
    backbone_model.base_model.encoder.parameters(),
    oneformer.model.pixel_level_module.encoder.encoder.parameters(),
)

```

## Who can review?

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.
@Cyrilvallez 

<!-- 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: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @zach-huggingface, @SunMarc and @qgallouedec
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @SunMarc @zach-huggingface
- 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: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Changed files:
filename: src/transformers/models/oneformer/modeling_oneformer.py
patch: @@ -2773,7 +2773,6 @@ def _init_weights(self, module: nn.Module):
         elif isinstance(module, OneFormerTransformerDecoder):
             nn.init.xavier_uniform_(module.query_input_projection.weight, gain=xavier_std)
             nn.init.constant_(module.query_input_projection.bias, 0)
-            module.query_input_projection._is_hf_initialized = True
         elif isinstance(module, OneFormerPixelDecoderEncoderMultiscaleDeformableAttention):
             nn.init.constant_(module.sampling_offsets.weight.data, 0.0)
             thetas = torch.arange(module.n_heads, dtype=torch.int64).float() * (2.0 * math.pi / module.n_heads)
@@ -2793,46 +2792,16 @@ def _init_weights(self, module: nn.Module):
             nn.init.constant_(module.value_proj.bias.data, 0.0)
             nn.init.xavier_uniform_(module.output_proj.weight.data)
             nn.init.constant_(module.output_proj.bias.data, 0.0)
-        elif isinstance(module, OneFormerPixelDecoderEncoderOnly):
-            for p in module.parameters():
-                if p.dim() > 1:
-                    nn.init.xavier_uniform_(p)
         elif isinstance(module, OneFormerPixelDecoder):
-            for p in module.parameters():
-                if p.dim() > 1:
-                    nn.init.xavier_uniform_(p)
             nn.init.normal_(module.level_embed, std=0)
-        elif isinstance(module, OneFormerTransformerDecoderSelfAttentionLayer):
-            for p in module.parameters():
-                if p.dim() > 1:
-                    nn.init.xavier_uniform_(p, gain=xavier_std)
-        elif isinstance(module, OneFormerTransformerDecoderCrossAttentionLayer):
-            for p in module.parameters():
-                if p.dim() > 1:
-                    nn.init.xavier_uniform_(p, gain=xavier_std)
-        elif isinstance(module, OneFormerTransformerDecoderFFNLayer):
+        elif isinstance(module, OneFormerTransformerDecoderLayer):
             for p in module.parameters():
                 if p.dim() > 1:
                     nn.init.xavier_uniform_(p, gain=xavier_std)
         elif isinstance(module, OneFormerTransformerDecoderQueryTransformer):
             for p in module.parameters():
                 if p.dim() > 1:
                     nn.init.xavier_uniform_(p, gain=xavier_std)
-        elif isinstance(module, OneFormerPixelLevelModule):
-            for submodule in module.modules():
-                if isinstance(submodule, (nn.Conv2d, nn.Linear)):
-                    submodule.weight.data.normal_(mean=0.0, std=std)
-                    if submodule.bias is not None:
-                        submodule.bias.data.zero_()
-        elif isinstance(module, OneFormerTextContextDecoder):
-            for submodule in module.modules():
-                if isinstance(submodule, nn.Linear):
-                    nn.init.trunc_normal_(submodule.weight, std=0.02)
-                    if isinstance(submodule, nn.Linear) and submodule.bias is not None:
-                        nn.init.constant_(submodule.bias, 0)
-                elif isinstance(submodule, nn.LayerNorm):
-                    nn.init.constant_(submodule.bias, 0)
-                    nn.init.constant_(submodule.weight, 1.0)
         elif isinstance(module, OneFormerTextTransformer):
             proj_std = (module.width**-0.5) * ((2 * module.num_layers) ** -0.5)
             attn_std = module.width**-0.5
@@ -2848,27 +2817,27 @@ def _init_weights(self, module: nn.Module):
         if hasattr(module, "reference_points"):
             nn.init.xavier_uniform_(module.reference_points.weight.data, gain=1.0)
             nn.init.constant_(module.reference_points.bias.data, 0.0)
-        elif isinstance(module, OneFormerTaskModel):
+        elif isinstance(module, OneFormerMLPPredictionHead):
             for submodule in module.modules():
-                if isinstance(module, OneFormerMLPPredictionHead):
-                    for submodule in module.modules():
-                        if isinstance(submodule, nn.Linear):
-                            nn.init.xavier_uniform_(submodule.weight, gain=xavier_std)
-                            nn.init.constant_(submodule.bias, 0)
-                        elif isinstance(module, nn.LayerNorm):
-                            module.bias.data.zero_()
-                            module.weight.data.fill_(1.0)
+                if isinstance(submodule, nn.Linear):
+                    nn.init.xavier_uniform_(submodule.weight, gain=xavier_std)
+                    nn.init.constant_(submodule.bias, 0)
         elif isinstance(module, nn.MultiheadAttention):
             module.in_proj_weight.data.normal_(mean=0.0, std=std)
             module.in_proj_bias.data.zero_()
         elif isinstance(module, (nn.Linear, nn.Conv2d, nn.BatchNorm2d)):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
                 module.bias.data.zero_()
+        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
         elif isinstance(module, nn.Embedding):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, OneFormerLoss):
+            module.logit_scale.data.fill_(np.log(1 / self.config.contrastive_temperature))
 
 
 @auto_docstring

filename: tests/models/oneformer/test_modeling_oneformer.py
patch: @@ -13,14 +13,13 @@
 # limitations under the License.
 """Testing suite for the PyTorch OneFormer model."""
 
-import copy
 import inspect
 import unittest
 
 import numpy as np
 
 from tests.test_modeling_common import floats_tensor
-from transformers import OneFormerConfig, is_torch_available, is_vision_available
+from transformers import AutoModelForImageClassification, OneFormerConfig, is_torch_available, is_vision_available
 from transformers.testing_utils import (
     is_flaky,
     require_timm,
@@ -35,7 +34,7 @@
 from transformers.utils import cached_property
 
 from ...test_configuration_common import ConfigTester
-from ...test_modeling_common import ModelTesterMixin
+from ...test_modeling_common import ModelTesterMixin, _config_zero_init
 from ...test_pipeline_mixin import PipelineTesterMixin
 
 
@@ -51,14 +50,6 @@
     from PIL import Image
 
 
-def _config_zero_init(config):
-    configs_no_init = copy.deepcopy(config)
-    for key in configs_no_init.__dict__.keys():
-        if "_range" in key or "_std" in key or "initializer_factor" in key or "layer_scale" in key:
-            setattr(configs_no_init, key, 1e-10)
-    return configs_no_init
-
-
 class OneFormerModelTester:
     def __init__(
         self,
@@ -375,19 +366,64 @@ def test_attention_outputs(self):
 
     def test_initialization(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        config.is_training = True
         config.contrastive_temperature = 1
 
         configs_no_init = _config_zero_init(config)
         for model_class in self.all_model_classes:
             model = model_class(config=configs_no_init)
             for name, param in model.named_parameters():
                 if param.requires_grad:
+                    if (
+                        "self_attn.sampling_offsets.bias" in name
+                        or "self_attn.value_proj.weight" in name
+                        or "self_attn.output_proj.weight" in name
+                        or "self_attn.in_proj_weight" in name
+                        or "self_attn.out_proj.weight" in name
+                        or "mlp.fc1.weight" in name
+                        or "mlp.fc2.weight" in name
+                        or "text_mapper.text_encoder.positional_embedding" in name
+                        or "text_mapper.text_encoder.token_embedding.weight" in name
+                    ):
+                        continue
                     self.assertIn(
                         ((param.data.mean() * 1e9).round() / 1e9).item(),
                         [0.0, 1.0],
                         msg=f"Parameter {name} of model {model_class} seems not properly initialized",
                     )
 
+    def test_initialization_pretrained_backbone(self):
+        backbone_name = "microsoft/resnet-18"
+
+        # load OneFormerConfig config with a pretrained backbone
+        config = OneFormerConfig(
+            backbone=backbone_name,
+            use_pretrained_backbone=True,
+        )
+
+        # load pretrained backbone
+        backbone_model = AutoModelForImageClassification.from_pretrained(backbone_name, device_map=torch_device)
+
+        def params_match(params1, params2):
+            return all((p1 == p2).all() for p1, p2 in zip(params1, params2))
+
+        for model_class in self.all_model_classes:
+            model = model_class(config).to(torch_device).eval()
+            if model.__class__.__name__ == "OneFormerModel":
+                self.assertTrue(
+                    params_match(
+                        backbone_model.base_model.encoder.parameters(),
+                        model.pixel_level_module.encoder.encoder.parameters(),
+                    )
+                )
+            elif model.__class__.__name__ == "OneFormerForUniversalSegmentation":
+                self.assertTrue(
+                    params_match(
+                        backbone_model.base_model.encoder.parameters(),
+                        model.model.pixel_level_module.encoder.encoder.parameters(),
+                    )
+                )
+
     def test_training(self):
         if not self.model_tester.is_training:
             self.skipTest(reason="model_tester.is_training is set to False")