Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit c9de109

Browse files
yonigozlanyonigozlan
authored
🚨 Refactor ViT to updated standards (#41693)
* refactor attention * use mlp instead of intermediate and output * use modular for deit * nit * fix deit (to check) and refactor beit * use modular for beit * style * refactor vit deit beit segformer * add pos_embed_utils, refactor swin * Fix after merge + refactor pixio * refactor most drop_path in transformers * refactor vit_msn * fix modular * refactor vivit and audio_spectrogram_transformer, remove obsolete copied from and add todos * Fix repo * fix nit * Fix drop_path_schedule * nit * use modular for pos embed * reverse DropPath as a util * add conversion ijepa + fix offload vit_mae * fix modular * fix qkv bias in output proj + swin attention per stage * address review comments * fix swinn * improve swin and beit * last fixes after review * standardize backbone beit * fix beit * fix beit * fix repo * changes after review * Fix beit * fix after review * fix unexpected key swin * Address review * fix repo * Fix after review * fix modular * fix depth pro * Fix conversion mappings * fix style * Fix failing tests * fix ci * Fix two slow failing tests
1 parent 381032b commit c9de109

106 files changed

Lines changed: 9314 additions & 5605 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎docs/source/en/auto_docstring.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,12 +89,12 @@ class MySpecialModel(PreTrainedModel):
8989
Also apply `@auto_docstring` to classes that inherit from [`~utils.ModelOutput`].
9090

9191
```python
92-
@dataclass
9392
@auto_docstring(
9493
custom_intro="""
9594
Custom model outputs with additional fields.
9695
"""
9796
)
97+
@dataclass
9898
class MyModelOutput(ImageClassifierOutput):
9999
r"""
100100
loss (`torch.FloatTensor`, *optional*):

‎src/transformers/conversion_mapping.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,11 +104,108 @@
104104
"MllamaModel": "LlavaModel",
105105
"MaskFormerDetrDecoder": "DetrModel",
106106
"Qwen2_5_VLForConditionalGeneration": "Qwen2VLForConditionalGeneration",
107+
# ViT-style vision models (old HuggingFace checkpoint format → new modular format)
108+
"ASTModel": "ViTModel",
109+
"BeitModel": "ViTModel",
110+
"DeiTModel": "ViTModel",
111+
"IJepaModel": "ViTModel",
112+
"ViTMAEModel": "ViTModel",
113+
"ViTMSNModel": "ViTModel",
114+
"VivitModel": "ViTModel",
107115
}
108116

109117

110118
def _build_checkpoint_conversion_mapping():
111119
mapping = {
120+
"ViTModel": [
121+
WeightRenaming(r"encoder\.layer\.", "layers."),
122+
WeightRenaming("attention.query", "q_proj"),
123+
WeightRenaming("attention.key", "k_proj"),
124+
WeightRenaming("attention.value", "v_proj"),
125+
WeightRenaming("attention.output.dense", "attention.o_proj"),
126+
WeightRenaming("intermediate.dense", "mlp.fc1"),
127+
WeightRenaming("output.dense", "mlp.fc2"),
128+
],
129+
"ViTMSNForImageClassification": [
130+
WeightRenaming(r"^encoder\.", "vit.encoder."),
131+
],
132+
"ViTMAEForPreTraining": [
133+
WeightRenaming("attention.query", "q_proj"),
134+
WeightRenaming("attention.key", "k_proj"),
135+
WeightRenaming("attention.value", "v_proj"),
136+
WeightRenaming("attention.output.dense", "attention.o_proj"),
137+
WeightRenaming("intermediate.dense", "mlp.fc1"),
138+
WeightRenaming("output.dense", "mlp.fc2"),
139+
],
140+
"BeitBackbone": [
141+
WeightRenaming(r"^fpn1\.0\.", "fpn.fpn1.conv_transpose1."),
142+
WeightRenaming(r"^fpn1\.1\.", "fpn.fpn1.normalization."),
143+
WeightRenaming(r"^fpn1\.3\.", "fpn.fpn1.conv_transpose2."),
144+
WeightRenaming(r"^fpn2\.0\.", "fpn.fpn2."),
145+
WeightRenaming(r"^encoder\.layer\.", "beit.encoder.layer."),
146+
WeightRenaming(r"^embeddings\.", "beit.embeddings."),
147+
WeightRenaming("attention.query", "q_proj"),
148+
WeightRenaming("attention.key", "k_proj"),
149+
WeightRenaming("attention.value", "v_proj"),
150+
WeightRenaming("attention.output.dense", "attention.o_proj"),
151+
WeightRenaming("intermediate.dense", "mlp.fc1"),
152+
WeightRenaming("output.dense", "mlp.fc2"),
153+
],
154+
"BeitForSemanticSegmentation": [
155+
WeightRenaming(r"(?<!psp_modules\.[0-9]\.1\.)bn\.", "normalization."),
156+
WeightRenaming(r"(?<!psp_modules\.[0-9]\.1\.)conv\.weight", "convolution.weight"),
157+
WeightRenaming(r"^fpn1\.0\.", "fpn.fpn1.conv_transpose1."),
158+
WeightRenaming(r"^fpn1\.1\.", "fpn.fpn1.normalization."),
159+
WeightRenaming(r"^fpn1\.3\.", "fpn.fpn1.conv_transpose2."),
160+
WeightRenaming(r"^fpn2\.0\.", "fpn.fpn2."),
161+
WeightRenaming("decode_head.bottleneck.", "decode_head.psp_bottleneck."),
162+
WeightRenaming(
163+
r"decode_head\.psp_modules\.(\d+)\.1\.conv\.weight",
164+
r"decode_head.psp_modules.blocks.\1.conv.convolution.weight",
165+
),
166+
WeightRenaming(
167+
r"decode_head\.psp_modules\.(\d+)\.1\.bn\.",
168+
r"decode_head.psp_modules.blocks.\1.conv.normalization.",
169+
),
170+
],
171+
"lw_detr": [
172+
WeightRenaming("attention.attention.query", "attention.q_proj"),
173+
WeightRenaming("attention.attention.key", "attention.k_proj"),
174+
WeightRenaming("attention.attention.value", "attention.v_proj"),
175+
WeightRenaming("attention.output", "attention.o_proj"),
176+
],
177+
"SegformerModel": [
178+
WeightRenaming(r"encoder.patch_embeddings.(\d+).", r"stages.\1.patch_embeddings."),
179+
WeightRenaming(r"encoder.block.(\d+).", r"stages.\1.blocks."),
180+
WeightRenaming(r"encoder.layer_norm.(\d+)", r"stages.\1.layer_norm"),
181+
WeightRenaming("attention.self.query", "attention.q_proj"),
182+
WeightRenaming("attention.self.key", "attention.k_proj"),
183+
WeightRenaming("attention.self.value", "attention.v_proj"),
184+
WeightRenaming("attention.self.sr", "attention.sequence_reduction.sequence_reduction"),
185+
WeightRenaming("attention.self.layer_norm", "attention.sequence_reduction.layer_norm"),
186+
WeightRenaming("attention.output.dense", "attention.o_proj"),
187+
WeightRenaming("mlp.dense1", "mlp.fc1"),
188+
WeightRenaming("mlp.dense2", "mlp.fc2"),
189+
WeightRenaming("layer_norm_1", "layernorm_before"),
190+
WeightRenaming("layer_norm_2", "layernorm_after"),
191+
],
192+
"SegformerForSemanticSegmentation": [WeightRenaming("decode_head.linear_c", "decode_head.linear_projections")],
193+
"swin": [
194+
WeightRenaming("attention.self.query", "attention.q_proj"),
195+
WeightRenaming("attention.self.key", "attention.k_proj"),
196+
WeightRenaming("attention.self.value", "attention.v_proj"),
197+
WeightRenaming(
198+
"attention.self.relative_position_bias_table",
199+
"attention.relative_position_bias.relative_position_bias_table",
200+
),
201+
WeightRenaming("attention.output.dense", "attention.o_proj"),
202+
WeightRenaming("intermediate.dense", "mlp.fc1"),
203+
WeightRenaming("output.dense", "mlp.fc2"),
204+
],
205+
"SwinBackbone": [
206+
WeightRenaming(r"^encoder\.", "swin.encoder."),
207+
WeightRenaming(r"^embeddings\.", "swin.embeddings."),
208+
],
112209
"altclip": [
113210
WeightRenaming(source_patterns=r"layer\.", target_patterns="layers."),
114211
],
@@ -548,6 +645,17 @@ def _build_checkpoint_conversion_mapping():
548645
operations=[ErnieFuseAndSplitTextVisionExperts(stack_dim=0, concat_dim=1)],
549646
),
550647
],
648+
# MaskFormer embeds both a Swin backbone (model_type="swin") and a DETR decoder
649+
# (model_type="detr") as submodules. Both get their mappings collected automatically, but
650+
# the Swin reverse mapping "mlp.fc1 → intermediate.dense" would corrupt DETR decoder keys
651+
# if it runs before the DETR reverse "layers.N.mlp.fc1 → layers.N.fc1".
652+
# By adding a "maskformer"-level mapping with the DETR fc1 rename, it is collected first
653+
# (model-level before submodule-level), so its reverse runs first and removes "mlp.fc1"
654+
# from DETR decoder paths before the Swin reverse can match them.
655+
"maskformer": [
656+
WeightRenaming(r"layers.(\d+).fc1", r"layers.\1.mlp.fc1"),
657+
WeightRenaming(r"layers.(\d+).fc2", r"layers.\1.mlp.fc2"),
658+
],
551659
"DetrModel": [
552660
WeightRenaming("backbone.conv_encoder", "backbone"),
553661
WeightRenaming("out_proj", "o_proj"),
@@ -836,6 +944,20 @@ def _build_checkpoint_conversion_mapping():
836944
mapping["ernie4_5_moe"] += [
837945
WeightRenaming("mlp.moe_statics.e_score_correction_bias", "mlp.gate.moe_statics.e_score_correction_bias")
838946
]
947+
mapping["BeitModel"] = mapping["ViTModel"].copy() + [
948+
WeightRenaming("encoder.relative_position_bias", "shared_position_bias"),
949+
WeightRenaming(
950+
"attention.attention.relative_position_bias.relative_position_bias_table",
951+
"relative_position_bias.relative_position_bias_table",
952+
),
953+
]
954+
955+
mapping["pixio"] = mapping["ViTModel"].copy()
956+
mapping["pixio"] += [
957+
WeightRenaming("norm1", "layernorm_before"),
958+
WeightRenaming("norm2", "layernorm_after"),
959+
WeightRenaming(r"^encoder\.", "pixio."),
960+
]
839961

840962
mapping["minimax_m2"] = mapping["mixtral"].copy()
841963
mapping["minimax_m2"] += [

‎src/transformers/loss/loss_utils.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,31 @@ def ForTokenClassification(logits: torch.Tensor, labels, config, **kwargs):
144144
return fixed_cross_entropy(logits, labels, **kwargs)
145145

146146

147+
def ForSemanticSegmentationLoss(
148+
logits: torch.Tensor,
149+
labels: torch.Tensor,
150+
ignore_index: int = 255,
151+
num_items_in_batch: torch.Tensor | None = None,
152+
auxiliary_logits: torch.Tensor | None = None,
153+
auxiliary_loss_weight: float = 0.4,
154+
**kwargs,
155+
) -> torch.Tensor:
156+
upsampled_logits = nn.functional.interpolate(logits, size=labels.shape[-2:], mode="bilinear", align_corners=False)
157+
loss = fixed_cross_entropy(
158+
upsampled_logits, labels, num_items_in_batch=num_items_in_batch, ignore_index=ignore_index
159+
)
160+
if auxiliary_logits is not None:
161+
upsampled_auxiliary_logits = nn.functional.interpolate(
162+
auxiliary_logits, size=labels.shape[-2:], mode="bilinear", align_corners=False
163+
)
164+
loss = loss + auxiliary_loss_weight * fixed_cross_entropy(
165+
upsampled_auxiliary_logits, labels, num_items_in_batch=num_items_in_batch, ignore_index=ignore_index
166+
)
167+
return loss
168+
169+
147170
LOSS_MAPPING = {
171+
"ForSemanticSegmentation": ForSemanticSegmentationLoss,
148172
"ForCausalLM": ForCausalLMLoss,
149173
"ForMaskedLM": ForMaskedLMLoss,
150174
"ForQuestionAnswering": ForQuestionAnsweringLoss,

‎src/transformers/models/aimv2/modeling_aimv2.py‎

Lines changed: 56 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,54 @@ def forward(self, x):
113113
return down_proj
114114

115115

116+
def build_2d_sinusoidal_position_embedding(
117+
height: int,
118+
width: int,
119+
embed_dim: int = 256,
120+
temperature: float = 10000.0,
121+
cls_token: bool = False,
122+
device: torch.device | None = None,
123+
dtype: torch.dtype = torch.float32,
124+
) -> torch.Tensor:
125+
"""2D sinusoidal position embeddings for an image patch grid.
126+
127+
Each (h, w) position gets an ``embed_dim``-dimensional vector laid out as
128+
``[sin_h | cos_h | sin_w | cos_w]``, with row-major (H-outer) patch ordering.
129+
130+
Args:
131+
height: Grid height in patches.
132+
width: Grid width in patches.
133+
embed_dim: Total embedding dimension; must be divisible by 4.
134+
temperature: Base for the frequency decay.
135+
cls_token: If `True`, prepend a zero row for a CLS token.
136+
device: Target device; defaults to CPU.
137+
dtype: Output dtype; frequency arithmetic uses float64 internally.
138+
139+
Returns:
140+
Tensor of shape ``(height * width [+1], embed_dim)``.
141+
"""
142+
if embed_dim % 4 != 0:
143+
raise ValueError(f"`embed_dim` must be divisible by 4, got {embed_dim}")
144+
145+
pos_dim = embed_dim // 4
146+
omega = torch.arange(pos_dim, dtype=torch.float64, device=device) / pos_dim
147+
omega = 1.0 / temperature**omega # (D/4,)
148+
149+
grid_h = torch.arange(height, dtype=torch.float64, device=device)
150+
grid_w = torch.arange(width, dtype=torch.float64, device=device)
151+
grid_h, grid_w = torch.meshgrid(grid_h, grid_w, indexing="ij") # (H, W) each
152+
153+
emb_h = grid_h.flatten().outer(omega) # (H*W, D/4)
154+
emb_w = grid_w.flatten().outer(omega) # (H*W, D/4)
155+
156+
pos_embed = torch.cat([emb_h.sin(), emb_h.cos(), emb_w.sin(), emb_w.cos()], dim=1)
157+
158+
if cls_token:
159+
pos_embed = torch.cat([torch.zeros(1, embed_dim, dtype=torch.float64, device=device), pos_embed], dim=0)
160+
161+
return pos_embed.to(dtype)
162+
163+
116164
class Aimv2VisionEmbeddings(nn.Module):
117165
def __init__(self, config: Aimv2VisionConfig):
118166
super().__init__()
@@ -128,36 +176,24 @@ def __init__(self, config: Aimv2VisionConfig):
128176
self.position_embedding = nn.Embedding(num_patches, config.hidden_size)
129177
self.register_buffer("position_ids", torch.arange(num_patches).expand((1, -1)), persistent=False)
130178

131-
@staticmethod
132-
def build_2d_sincos_position_embedding(
133-
height, width, embed_dim=256, temperature=10000.0, device="cpu", dtype=torch.float32
134-
) -> torch.Tensor:
135-
grid_w = torch.arange(int(width), dtype=dtype, device=device)
136-
grid_h = torch.arange(int(height), dtype=dtype, device=device)
137-
grid_h, grid_w = torch.meshgrid(grid_w, grid_h, indexing="xy")
138-
139-
pos_dim = embed_dim // 4
140-
omega = torch.arange(pos_dim, dtype=dtype, device=device) / pos_dim
141-
omega = 1.0 / (temperature**omega)
142-
143-
out_h = grid_h.flatten()[..., None] @ omega[None, :]
144-
out_w = grid_w.flatten()[..., None] @ omega[None, :]
145-
146-
return torch.concat([out_h.sin(), out_h.cos(), out_w.sin(), out_w.cos()], dim=1)[None, :, :]
147-
148179
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
149180
_, _, height, width = pixel_values.size()
150181
hidden_states = self.patch_embed(pixel_values).flatten(2).transpose(1, 2)
151182
hidden_states = self.rms_norm(hidden_states)
152183

153184
if self.config.is_native:
154-
pos_embed = self.build_2d_sincos_position_embedding(
155-
height // self.patch_size,
156-
width // self.patch_size,
185+
pos_embed = build_2d_sinusoidal_position_embedding(
186+
height=height // self.patch_size,
187+
width=width // self.patch_size,
157188
embed_dim=self.config.hidden_size,
158189
device=hidden_states.device,
159190
dtype=hidden_states.dtype,
160191
)
192+
# AIMv2 was trained with [sin_w|cos_w|sin_h|cos_h] layout (matching ViT-MAE's
193+
# original naming-bug convention); rotate the canonical h-first embedding to match.
194+
half = pos_embed.shape[-1] // 2
195+
pos_embed = torch.cat([pos_embed[..., half:], pos_embed[..., :half]], dim=-1)
196+
pos_embed = pos_embed.unsqueeze(0)
161197
else:
162198
pos_embed = self.position_embedding(self.position_ids)
163199

‎src/transformers/models/aimv2/modular_aimv2.py‎

Lines changed: 9 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
from ..llama.modeling_llama import LlamaMLP, LlamaRMSNorm
3636
from ..siglip.configuration_siglip import SiglipConfig, SiglipTextConfig, SiglipVisionConfig
3737
from ..siglip.modeling_siglip import SiglipAttention, SiglipEncoder, SiglipOutput
38+
from ..vit_mae.modeling_vit_mae import build_2d_sinusoidal_position_embedding
3839

3940

4041
@auto_docstring(checkpoint="apple/aimv2-large-patch14-224-lit")
@@ -164,36 +165,24 @@ def __init__(self, config: Aimv2VisionConfig):
164165
self.position_embedding = nn.Embedding(num_patches, config.hidden_size)
165166
self.register_buffer("position_ids", torch.arange(num_patches).expand((1, -1)), persistent=False)
166167

167-
@staticmethod
168-
def build_2d_sincos_position_embedding(
169-
height, width, embed_dim=256, temperature=10000.0, device="cpu", dtype=torch.float32
170-
) -> torch.Tensor:
171-
grid_w = torch.arange(int(width), dtype=dtype, device=device)
172-
grid_h = torch.arange(int(height), dtype=dtype, device=device)
173-
grid_h, grid_w = torch.meshgrid(grid_w, grid_h, indexing="xy")
174-
175-
pos_dim = embed_dim // 4
176-
omega = torch.arange(pos_dim, dtype=dtype, device=device) / pos_dim
177-
omega = 1.0 / (temperature**omega)
178-
179-
out_h = grid_h.flatten()[..., None] @ omega[None, :]
180-
out_w = grid_w.flatten()[..., None] @ omega[None, :]
181-
182-
return torch.concat([out_h.sin(), out_h.cos(), out_w.sin(), out_w.cos()], dim=1)[None, :, :]
183-
184168
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
185169
_, _, height, width = pixel_values.size()
186170
hidden_states = self.patch_embed(pixel_values).flatten(2).transpose(1, 2)
187171
hidden_states = self.rms_norm(hidden_states)
188172

189173
if self.config.is_native:
190-
pos_embed = self.build_2d_sincos_position_embedding(
191-
height // self.patch_size,
192-
width // self.patch_size,
174+
pos_embed = build_2d_sinusoidal_position_embedding(
175+
height=height // self.patch_size,
176+
width=width // self.patch_size,
193177
embed_dim=self.config.hidden_size,
194178
device=hidden_states.device,
195179
dtype=hidden_states.dtype,
196180
)
181+
# AIMv2 was trained with [sin_w|cos_w|sin_h|cos_h] layout (matching ViT-MAE's
182+
# original naming-bug convention); rotate the canonical h-first embedding to match.
183+
half = pos_embed.shape[-1] // 2
184+
pos_embed = torch.cat([pos_embed[..., half:], pos_embed[..., :half]], dim=-1)
185+
pos_embed = pos_embed.unsqueeze(0)
197186
else:
198187
pos_embed = self.position_embedding(self.position_ids)
199188

0 commit comments

Comments
 (0)