Support "MultiStepLR" and "StepLR" learning rate schedulers.

ml4chem/optim/handler.py

@@ -147,10 +147,36 @@ def get_lr_scheduler(optimizer, lr_scheduler):
         scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **kwargs)
         name = "ReduceLROnPlateau"
 
+    elif scheduler_name == "multisteplr":
+        scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, **kwargs)
+        name = "MultiStepLR"
+
+    elif scheduler_name == "steplr":
+        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, **kwargs)
+        name = "StepLR"
+
+
     logger.info("Learning Rate Scheduler")
     logger.info("-----------------------")
     logger.info("    - Name: {}.".format(name))
     logger.info("    - Args: {}.".format(kwargs))
     logger.info("")
 
     return scheduler
+
+
+def get_lr(optimizer):
+    """Get current learning rate
+    
+    Parameters
+    ----------
+    optimizer : obj
+        An optimizer object.
+    
+    Returns
+    -------
+    lr
+        Current learning rate.
+    """
+    for param_group in optimizer.param_groups:
+        return param_group['lr']

ml4chem/potentials.py

@@ -216,6 +216,7 @@ def train(
         lossfxn=None,
         regularization=0.0,
         batch_size=None,
+        **kwargs
     ):
         """Method to train models
 
@@ -315,6 +316,7 @@ def train(
                 lossfxn=lossfxn,
                 device=device,
                 batch_size=batch_size,
+                **kwargs
             )
 
         self.save(self.model, features=self.features, path=self.path, label=self.label)