Fix beit init

Author: Niels Rogge

src/transformers/models/beit/configuration_beit.py

@@ -59,6 +59,8 @@ class BeitConfig(PretrainedConfig):
  The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
  attention_probs_dropout_prob (`float`, *optional*, defaults to 0.0):
  The dropout ratio for the attention probabilities.
+ cls_token_initializer_range (`float`, *optional*, defaults to 0.02):
+ The standard deviation of the truncated_normal_initializer for initializing the `cls_token` parameter.
  initializer_range (`float`, *optional*, defaults to 0.02):
  The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
  layer_norm_eps (`float`, *optional*, defaults to 1e-12):
@@ -127,6 +129,7 @@ def __init__(
  hidden_act="gelu",
  hidden_dropout_prob=0.0,
  attention_probs_dropout_prob=0.0,
+ cls_token_initializer_range=0.02,
  initializer_range=0.02,
  layer_norm_eps=1e-12,
  image_size=224,
@@ -159,6 +162,7 @@ def __init__(
  self.hidden_act = hidden_act
  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.initializer_range = initializer_range
  self.layer_norm_eps = layer_norm_eps
 

src/transformers/models/beit/modeling_beit.py

@@ -560,19 +560,17 @@ class BeitPreTrainedModel(PreTrainedModel):
 
  def _init_weights(self, module):
  """Initialize the weights"""
- if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
- # Slightly different from the TF version which uses truncated_normal for initialization
- # cf https://github.com/pytorch/pytorch/pull/5617
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.bias is not None:
- module.bias.data.zero_()
- elif isinstance(module, nn.Embedding):
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
- if module.padding_idx is not None:
- module.weight.data[module.padding_idx].zero_()
+ if isinstance(module, nn.Linear):
+ torch.nn.init.trunc_normal_(module.data, std=self.config.initializer_range)
  elif isinstance(module, nn.LayerNorm):
  module.bias.data.zero_()
  module.weight.data.fill_(1.0)
+ elif isinstance(module, BeitEmbeddings):
+ module.cls_token.data.normal_(mean=0.0, std=self.config.cls_token_initializer_range)
+ if module.mask_token is not None:
+ torch.nn.init.trunc_normal_(module.mask_token.data, std=self.config.initializer_range)
+ if module.position_embeddings is not None:
+ torch.nn.init.trunc_normal_(module.position_embeddings.data, std=self.config.initializer_range)
 
  def _set_gradient_checkpointing(self, module, value=False):
  if isinstance(module, BeitEncoder):