Improve initialization

Author: Niels Rogge

src/transformers/models/deit/configuration_deit.py

@@ -66,6 +66,9 @@ class DeiTConfig(PretrainedConfig):
  The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
  cls_token_initializer_range (`float`, *optional*, defaults to 1e-6):
  The standard deviation of the truncated_normal_initializer for initializing the `cls_token` parameter.
+ distillation_token_initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing the `distillation_token`
+ parameter.
  layer_norm_eps (`float`, *optional*, defaults to 1e-12):
  The epsilon used by the layer normalization layers.
  image_size (`int`, *optional*, defaults to `224`):
@@ -95,7 +98,6 @@ class DeiTConfig(PretrainedConfig):
  ```"""
  model_type = "deit"
 
- # Copied from transformers.models.vit.configuration_vit.ViTConfig.__init__
  def __init__(
  self,
  hidden_size=768,
@@ -106,6 +108,7 @@ def __init__(
  hidden_dropout_prob=0.0,
  attention_probs_dropout_prob=0.0,
  cls_token_initializer_range=1e-6,
+ distillation_token_initializer_range=0.02,
  initializer_range=0.02,
  layer_norm_eps=1e-12,
  image_size=224,
@@ -125,6 +128,7 @@ def __init__(
  self.hidden_dropout_prob = hidden_dropout_prob
  self.attention_probs_dropout_prob = attention_probs_dropout_prob
  self.cls_token_initializer_range = cls_token_initializer_range
+ self.distillation_token_initializer_range = distillation_token_initializer_range
  self.initializer_range = initializer_range
  self.layer_norm_eps = layer_norm_eps
  self.image_size = image_size

src/transformers/models/deit/modeling_deit.py

@@ -387,7 +387,6 @@ def custom_forward(*inputs):
  )
 
 
-# Copied from transformers.models.vit.modeling_vit.ViTPreTrainedModel with ViT->DeiT all-casing
 class DeiTPreTrainedModel(PreTrainedModel):
  """
  An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
@@ -415,6 +414,9 @@ def _init_weights(self, module) -> None:
  module.weight.data.fill_(1.0)
  elif isinstance(module, DeiTEmbeddings):
  module.cls_token.data.normal_(mean=0.0, std=self.config.cls_token_initializer_range)
+ torch.nn.init.trunc_normal_(
+ module.distillation_token.data, std=self.config.distillation_token_initializer_range
+ )
  if module.mask_token is not None:
  torch.nn.init.trunc_normal_(module.mask_token.data, std=self.config.initializer_range)
  torch.nn.init.trunc_normal_(module.position_embeddings.data, std=self.config.initializer_range)