Repository: huggingface/transformers
Pull request: [LW-DETR] Fix training

# What does this PR do?

This PR fixes the training of LW-DETR. It turned out that the model was not able to overfit a single batch. Hence I asked Codex to investigate this. 

It turns out there were 3 bugs:

1. A logits calibration gap, caused by a training-time bug: the encoder output box coordinates are incorrectly detached in the HF model, cutting off gradient flow from the encoder output box loss to the two-stage bbox prediction heads.
2. Train-time group-DETR self-attention in HF used a reshape based on pre-split `(batch, seq_len)` instead of post-split shape.
3. (nit) Converted HF configs were using `dropout=0.1`, while original LW-DETR training config uses `dropout=0.0`.

After this, I was able to fine-tune LW-DETR properly on a sample dataset (https://huggingface.co/datasets/nielsr/tray-cart-detection) with parity with respect to the original training repo: https://github.com/Atten4Vis/LW-DETR.

cc @sbucaille 


Changed files:
filename: src/transformers/models/lw_detr/configuration_lw_detr.py
patch: @@ -179,7 +179,7 @@ def __init__(
         batch_norm_eps=1e-5,
         # decoder
         d_model=256,
-        dropout=0.1,
+        dropout=0.0,
         decoder_ffn_dim=2048,
         decoder_n_points=4,
         decoder_layers: int = 3,

filename: src/transformers/models/lw_detr/modeling_lw_detr.py
patch: @@ -672,8 +672,6 @@ def forward(
         **kwargs: Unpack[TransformersKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, seq_len, _ = hidden_states.shape
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
 
         hidden_states_original = hidden_states
         if position_embeddings is not None:
@@ -687,6 +685,8 @@ def forward(
             )
             hidden_states = torch.cat(hidden_states.split(seq_len // self.config.group_detr, dim=1), dim=0)
 
+        attention_input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*attention_input_shape, -1, self.head_dim)
         query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
         key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
         value_states = self.v_proj(hidden_states_original).view(hidden_shape).transpose(1, 2)
@@ -705,7 +705,7 @@ def forward(
             scaling=self.scaling,
             **kwargs,
         )
-        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = attn_output.reshape(*attention_input_shape, -1).contiguous()
         attn_output = self.o_proj(attn_output)
 
         if self.training:
@@ -1261,26 +1261,6 @@ def get_valid_ratio(self, mask, dtype=torch.float32):
         valid_ratio = torch.stack([valid_ratio_width, valid_ratio_height], -1)
         return valid_ratio
 
-    def get_proposal_pos_embed(self, proposals):
-        """Get the position embedding of the proposals."""
-
-        num_pos_feats = self.config.d_model // 2
-        temperature = 10000
-        scale = 2 * math.pi
-
-        # Compute position embeddings in float32 to avoid overflow with large temperature values in fp16
-        proposals_dtype = proposals.dtype
-        dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=proposals.device)
-        dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats)
-        # batch_size, num_queries, 4
-        proposals = proposals.sigmoid().to(torch.float32) * scale
-        # batch_size, num_queries, 4, 128
-        pos = proposals[:, :, :, None] / dim_t
-        # batch_size, num_queries, 4, 64, 2 -> batch_size, num_queries, 512
-        pos = torch.stack((pos[:, :, :, 0::2].sin(), pos[:, :, :, 1::2].cos()), dim=4).flatten(2)
-        # Convert back to target dtype after all computations are done
-        return pos.to(proposals_dtype)
-
     def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes):
         """Generate the encoder output proposals from encoded enc_output.
 
@@ -1293,8 +1273,10 @@ def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes)
             `tuple(torch.FloatTensor)`: A tuple of feature map and bbox prediction.
                 - object_query (Tensor[batch_size, sequence_length, hidden_size]): Object query features. Later used to
                   directly predict a bounding box. (without the need of a decoder)
-                - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals, after an inverse
-                  sigmoid.
+                - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals in [0, 1] space.
+                  Invalid positions (padding or out-of-bounds) are filled with 0.
+                - invalid_mask (Tensor[batch_size, sequence_length, 1]): Boolean mask that is True for invalid positions
+                  (padded pixels or proposals whose coordinates fall outside (0.01, 0.99)).
         """
         batch_size = enc_output.shape[0]
         proposals = []
@@ -1331,14 +1313,14 @@ def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes)
             _cur += height * width
         output_proposals = torch.cat(proposals, 1)
         output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True)
-        output_proposals = output_proposals.masked_fill(padding_mask.unsqueeze(-1), float("inf"))
-        output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf"))
+        invalid_mask = padding_mask | ~output_proposals_valid.squeeze(-1)
+        invalid_mask = padding_mask.unsqueeze(-1) | ~output_proposals_valid
+        output_proposals = output_proposals.masked_fill(invalid_mask, float(0))
 
         # assign each pixel as an object query
         object_query = enc_output
-        object_query = object_query.masked_fill(padding_mask.unsqueeze(-1), float(0))
-        object_query = object_query.masked_fill(~output_proposals_valid, float(0))
-        return object_query, output_proposals
+        object_query = object_query.masked_fill(invalid_mask, float(0))
+        return object_query, output_proposals, invalid_mask
 
     @can_return_tuple
     @auto_docstring
@@ -1421,7 +1403,7 @@ def forward(
         target = query_feat.unsqueeze(0).expand(batch_size, -1, -1)
         reference_points = reference_points.unsqueeze(0).expand(batch_size, -1, -1)
 
-        object_query_embedding, output_proposals = self.gen_encoder_output_proposals(
+        object_query_embedding, output_proposals, invalid_mask = self.gen_encoder_output_proposals(
             source_flatten, ~mask_flatten, spatial_shapes_list
         )
 
@@ -1436,6 +1418,7 @@ def forward(
             group_object_query = self.enc_output_norm[group_id](group_object_query)
 
             group_enc_outputs_class = self.enc_out_class_embed[group_id](group_object_query)
+            group_enc_outputs_class = group_enc_outputs_class.masked_fill(invalid_mask, float("-inf"))
             group_delta_bbox = self.enc_out_bbox_embed[group_id](group_object_query)
             group_enc_outputs_coord = refine_bboxes(output_proposals, group_delta_bbox)
 
@@ -1459,9 +1442,9 @@ def forward(
         object_query_undetach = torch.cat(object_query_undetach, 1)
 
         enc_outputs_class = object_query_undetach
-        enc_outputs_coord_logits = topk_coords_logits
+        enc_outputs_coord_logits = topk_coords_logits_undetach
 
-        reference_points = refine_bboxes(topk_coords_logits_undetach, reference_points)
+        reference_points = refine_bboxes(topk_coords_logits, reference_points)
 
         init_reference_points = reference_points
         decoder_outputs = self.decoder(

filename: src/transformers/models/lw_detr/modular_lw_detr.py
patch: @@ -213,7 +213,7 @@ def __init__(
         batch_norm_eps=1e-5,
         # decoder
         d_model=256,
-        dropout=0.1,
+        dropout=0.0,
         decoder_ffn_dim=2048,
         decoder_n_points=4,
         decoder_layers: int = 3,
@@ -733,8 +733,6 @@ def forward(
         **kwargs: Unpack[TransformersKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, seq_len, _ = hidden_states.shape
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, self.head_dim)
 
         hidden_states_original = hidden_states
         if position_embeddings is not None:
@@ -748,6 +746,8 @@ def forward(
             )
             hidden_states = torch.cat(hidden_states.split(seq_len // self.config.group_detr, dim=1), dim=0)
 
+        attention_input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*attention_input_shape, -1, self.head_dim)
         query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
         key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
         value_states = self.v_proj(hidden_states_original).view(hidden_shape).transpose(1, 2)
@@ -766,7 +766,7 @@ def forward(
             scaling=self.scaling,
             **kwargs,
         )
-        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = attn_output.reshape(*attention_input_shape, -1).contiguous()
         attn_output = self.o_proj(attn_output)
 
         if self.training:
@@ -1128,8 +1128,10 @@ def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes)
             `tuple(torch.FloatTensor)`: A tuple of feature map and bbox prediction.
                 - object_query (Tensor[batch_size, sequence_length, hidden_size]): Object query features. Later used to
                   directly predict a bounding box. (without the need of a decoder)
-                - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals, after an inverse
-                  sigmoid.
+                - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals in [0, 1] space.
+                  Invalid positions (padding or out-of-bounds) are filled with 0.
+                - invalid_mask (Tensor[batch_size, sequence_length, 1]): Boolean mask that is True for invalid positions
+                  (padded pixels or proposals whose coordinates fall outside (0.01, 0.99)).
         """
         batch_size = enc_output.shape[0]
         proposals = []
@@ -1166,14 +1168,14 @@ def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes)
             _cur += height * width
         output_proposals = torch.cat(proposals, 1)
         output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True)
-        output_proposals = output_proposals.masked_fill(padding_mask.unsqueeze(-1), float("inf"))
-        output_proposals = output_proposals.masked_fill(~output_proposals_valid, float("inf"))
+        invalid_mask = padding_mask | ~output_proposals_valid.squeeze(-1)
+        invalid_mask = padding_mask.unsqueeze(-1) | ~output_proposals_valid
+        output_proposals = output_proposals.masked_fill(invalid_mask, float(0))
 
         # assign each pixel as an object query
         object_query = enc_output
-        object_query = object_query.masked_fill(padding_mask.unsqueeze(-1), float(0))
-        object_query = object_query.masked_fill(~output_proposals_valid, float(0))
-        return object_query, output_proposals
+        object_query = object_query.masked_fill(invalid_mask, float(0))
+        return object_query, output_proposals, invalid_mask
 
     @can_return_tuple
     @auto_docstring
@@ -1256,7 +1258,7 @@ def forward(
         target = query_feat.unsqueeze(0).expand(batch_size, -1, -1)
         reference_points = reference_points.unsqueeze(0).expand(batch_size, -1, -1)
 
-        object_query_embedding, output_proposals = self.gen_encoder_output_proposals(
+        object_query_embedding, output_proposals, invalid_mask = self.gen_encoder_output_proposals(
             source_flatten, ~mask_flatten, spatial_shapes_list
         )
 
@@ -1271,6 +1273,7 @@ def forward(
             group_object_query = self.enc_output_norm[group_id](group_object_query)
 
             group_enc_outputs_class = self.enc_out_class_embed[group_id](group_object_query)
+            group_enc_outputs_class = group_enc_outputs_class.masked_fill(invalid_mask, float("-inf"))
             group_delta_bbox = self.enc_out_bbox_embed[group_id](group_object_query)
             group_enc_outputs_coord = refine_bboxes(output_proposals, group_delta_bbox)
 
@@ -1294,9 +1297,9 @@ def forward(
         object_query_undetach = torch.cat(object_query_undetach, 1)
 
         enc_outputs_class = object_query_undetach
-        enc_outputs_coord_logits = topk_coords_logits
+        enc_outputs_coord_logits = topk_coords_logits_undetach
 
-        reference_points = refine_bboxes(topk_coords_logits_undetach, reference_points)
+        reference_points = refine_bboxes(topk_coords_logits, reference_points)
 
         init_reference_points = reference_points
         decoder_outputs = self.decoder(
@@ -1323,6 +1326,9 @@ def forward(
             cross_attentions=decoder_outputs.cross_attentions,
         )
 
+    def get_proposal_pos_embed(self, proposals):
+        raise NotImplementedError("get_proposal_pos_embed is not used in LwDetrForObjectDetection")
+
 
 class LwDetrMLPPredictionHead(DeformableDetrMLPPredictionHead):
     pass