Adding optim schedule feature for codertimo#17

Author: codertimo

bert_pytorch/trainer/optim_schedule.py

@@ -0,0 +1,35 @@
+'''A wrapper class for optimizer '''
+import numpy as np
+
+
+class ScheduledOptim():
+    '''A simple wrapper class for learning rate scheduling'''
+
+    def __init__(self, optimizer, d_model, n_warmup_steps):
+        self._optimizer = optimizer
+        self.n_warmup_steps = n_warmup_steps
+        self.n_current_steps = 0
+        self.init_lr = np.power(d_model, -0.5)
+
+    def step_and_update_lr(self):
+        "Step with the inner optimizer"
+        self._update_learning_rate()
+        self._optimizer.step()
+
+    def zero_grad(self):
+        "Zero out the gradients by the inner optimizer"
+        self._optimizer.zero_grad()
+
+    def _get_lr_scale(self):
+        return np.min([
+            np.power(self.n_current_steps, -0.5),
+            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])
+
+    def _update_learning_rate(self):
+        ''' Learning rate scheduling per step '''
+
+        self.n_current_steps += 1
+        lr = self.init_lr * self._get_lr_scale()
+
+        for param_group in self._optimizer.param_groups:
+            param_group['lr'] = lr

bert_pytorch/trainer/pretrain.py

@@ -4,6 +4,7 @@
 from torch.utils.data import DataLoader
 
 from ..model import BERTLM, BERT
+from .optim_schedule import ScheduledOptim
 
 import tqdm
 
@@ -21,7 +22,7 @@ class BERTTrainer:
 
     def __init__(self, bert: BERT, vocab_size: int,
                  train_dataloader: DataLoader, test_dataloader: DataLoader = None,
-                 lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01,
+                 lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.01, warmup_steps=10000,
                  with_cuda: bool = True, cuda_devices=None, log_freq: int = 10):
         """
         :param bert: BERT model which you want to train
@@ -55,6 +56,7 @@ def __init__(self, bert: BERT, vocab_size: int,
 
         # Setting the Adam optimizer with hyper-param
         self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)
+        self.optim_schedule = ScheduledOptim(self.optim, self.bert.hidden, n_warmup_steps=warmup_steps)
 
         # Using Negative Log Likelihood Loss function for predicting the masked_token
         self.criterion = nn.NLLLoss(ignore_index=0)
@@ -110,9 +112,9 @@ def iteration(self, epoch, data_loader, train=True):
 
             # 3. backward and optimization only in train
             if train:
-                self.optim.zero_grad()
+                self.optim_schedule.zero_grad()
                 loss.backward()
-                self.optim.step()
+                self.optim_schedule.step_and_update_lr()
 
             # next sentence prediction accuracy
             correct = next_sent_output.argmax(dim=-1).eq(data["is_next"]).sum().item()