fix - update bert model weight init

ne7ermore · Aug 8, 2019 · 17caad1 · 17caad1
1 parent 254c133
commit 17caad1
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/BERT/model.py b/BERT/model.py
@@ -142,10 +142,10 @@ def reshape(x):
         return self.lm(outputs + residual)
 
     def reset_parameters(self):
-        self.w_qs.data.normal_(INIT_RANGE)
-        self.w_ks.data.normal_(INIT_RANGE)
-        self.w_vs.data.normal_(INIT_RANGE)
-        self.w_o.weight.data.normal_(INIT_RANGE)
+        self.w_qs.data.normal_(std=INIT_RANGE)
+        self.w_ks.data.normal_(std=INIT_RANGE)
+        self.w_vs.data.normal_(std=INIT_RANGE)
+        self.w_o.weight.data.normal_(std=INIT_RANGE)
 
 
 class EncoderLayer(nn.Module):
@@ -166,7 +166,7 @@ def __init__(self, d_model):
         super().__init__()
 
         self.linear = nn.Linear(d_model, d_model)
-        self.linear.weight.data.normal_(INIT_RANGE)
+        self.linear.weight.data.normal_(std=INIT_RANGE)
         self.linear.bias.data.zero_()
 
     def forward(self, x):
@@ -199,10 +199,10 @@ def __init__(self, args):
         self.gelu = GELU()
 
     def reset_parameters(self):
-        self.enc_ebd.weight.data.normal_(INIT_RANGE)
-        self.seg_ebd.weight.data.normal_(INIT_RANGE)
+        self.enc_ebd.weight.data.normal_(std=INIT_RANGE)
+        self.seg_ebd.weight.data.normal_(std=INIT_RANGE)
 
-        self.transform.weight.data.normal_(INIT_RANGE)
+        self.transform.weight.data.normal_(std=INIT_RANGE)
         self.transform.bias.data.zero_()
 
     def forward(self, inp, pos, segment_label):