Fix drop path being ignored in DINOv2 (#29147)

Fix drop path not being used
This commit is contained in:
Fernando Pérez-García 2024-02-20 17:31:59 +00:00 committed by GitHub
parent 3c00b885b9
commit c47576ca6e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -380,7 +380,7 @@ class Dinov2Layer(nn.Module):
self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.attention = Dinov2Attention(config)
self.layer_scale1 = Dinov2LayerScale(config)
self.drop_path1 = Dinov2DropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
self.drop_path = Dinov2DropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
@ -389,7 +389,6 @@ class Dinov2Layer(nn.Module):
else:
self.mlp = Dinov2MLP(config)
self.layer_scale2 = Dinov2LayerScale(config)
self.drop_path2 = Dinov2DropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
def forward(
self,
@ -408,7 +407,7 @@ class Dinov2Layer(nn.Module):
outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
# first residual connection
hidden_states = attention_output + hidden_states
hidden_states = self.drop_path(attention_output) + hidden_states
# in Dinov2, layernorm is also applied after self-attention
layer_output = self.norm2(hidden_states)
@ -416,7 +415,7 @@ class Dinov2Layer(nn.Module):
layer_output = self.layer_scale2(layer_output)
# second residual connection
layer_output = layer_output + hidden_states
layer_output = self.drop_path(layer_output) + hidden_states
outputs = (layer_output,) + outputs