Update modeling_norbert.py

modeling_norbert.py

@@ -59,13 +59,6 @@ class MaskClassifier(nn.Module):
         )
         self.initialize(config.hidden_size, subword_embedding)
 
-    def initialize(self, hidden_size, embedding):
-        std = math.sqrt(2.0 / (5.0 * hidden_size))
-        nn.init.trunc_normal_(self.nonlinearity[1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        self.nonlinearity[-1].weight = embedding
-        self.nonlinearity[1].bias.data.zero_()
-        self.nonlinearity[-1].bias.data.zero_()
-
     def forward(self, x, masked_lm_labels=None):
         if masked_lm_labels is not None:
             x = torch.index_select(x.flatten(0, 1), 0, torch.nonzero(masked_lm_labels.flatten() != -100).squeeze())
@@ -106,11 +99,6 @@ class FeedForward(nn.Module):
         )
         self.initialize(config.hidden_size)
 
-    def initialize(self, hidden_size):
-        std = math.sqrt(2.0 / (5.0 * hidden_size))
-        nn.init.trunc_normal_(self.mlp[1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        nn.init.trunc_normal_(self.mlp[-2].weight, mean=0.0, std=std, a=-2*std, b=2*std)
-
     def forward(self, x):
         return self.mlp(x)
 
@@ -170,15 +158,6 @@ class Attention(nn.Module):
         bucket_pos = torch.where(abs_pos <= mid, relative_pos, log_pos * sign).long()
         return bucket_pos
 
-    def initialize(self):
-        std = math.sqrt(2.0 / (5.0 * self.hidden_size))
-        nn.init.trunc_normal_(self.in_proj_qk.weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        nn.init.trunc_normal_(self.in_proj_v.weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        nn.init.trunc_normal_(self.out_proj.weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        self.in_proj_qk.bias.data.zero_()
-        self.in_proj_v.bias.data.zero_()
-        self.out_proj.bias.data.zero_()
-
     def compute_attention_scores(self, hidden_states, relative_embedding):
         key_len, batch_size, _ = hidden_states.size()
         query_len = key_len
@@ -246,13 +225,6 @@ class Embedding(nn.Module):
         self.relative_embedding = nn.Parameter(torch.empty(2 * config.position_bucket_size - 1, config.hidden_size))
         self.relative_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
 
-        self.initialize()
-
-    def initialize(self):
-        std = math.sqrt(2.0 / (5.0 * self.hidden_size))
-        nn.init.trunc_normal_(self.relative_embedding, mean=0.0, std=std, a=-2*std, b=2*std)
-        nn.init.trunc_normal_(self.word_embedding.weight, mean=0.0, std=std, a=-2*std, b=2*std)
-
     def forward(self, input_ids):
         word_embedding = self.dropout(self.word_layer_norm(self.word_embedding(input_ids)))
         relative_embeddings = self.relative_layer_norm(self.relative_embedding)
@@ -273,7 +245,18 @@ class NorbertPreTrainedModel(PreTrainedModel):
             module.activation_checkpointing = value
 
     def _init_weights(self, module):
-        pass  # everything is already initialized
+        std = math.sqrt(2.0 / (5.0 * self.hidden_size))
+
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight.data, mean=0.0, std=std, a=-2*std, b=2*std)
+            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):
+            nn.init.trunc_normal_(module.weight.data, mean=0.0, std=std, a=-2*std, b=2*std)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
 
 
 class NorbertModel(NorbertPreTrainedModel):
@@ -414,15 +397,6 @@ class Classifier(nn.Module):
             nn.Dropout(drop_out),
             nn.Linear(config.hidden_size, num_labels)
         )
-        self.hidden_size = config.hidden_size
-        self._init_weights()
-
-    def _init_weights(self):
-        std = math.sqrt(2.0 / (5.0 * self.hidden_size))
-        nn.init.trunc_normal_(self.nonlinearity[1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        nn.init.trunc_normal_(self.nonlinearity[-1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
-        self.nonlinearity[1].bias.data.zero_()
-        self.nonlinearity[-1].bias.data.zero_()
 
     def forward(self, x):
         x = self.nonlinearity(x)
@@ -439,12 +413,6 @@ class NorbertForSequenceClassification(NorbertModel):
         self.num_labels = config.num_labels
         self.head = Classifier(config, self.num_labels)
 
-    def post_init(self):
-        self.head._init_weights()
-
-    def _init_weights(self):
-        self.head._init_weights()
-
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,
@@ -511,9 +479,6 @@ class NorbertForTokenClassification(NorbertModel):
         self.num_labels = config.num_labels
         self.head = Classifier(config, self.num_labels)
 
-    def post_init(self):
-        self.head._init_weights()
-
     def forward(
         self,
         input_ids: Optional[torch.Tensor] = None,