Backend PyTorch supports "Step" LR decay (#722)

Author: Jerry-Jzy

deepxde/model.py

@@ -47,6 +47,8 @@ def __init__(self, data, net):
         if backend_name == "tensorflow.compat.v1":
             self.sess = None
             self.saver = None
+        elif backend_name == "pytorch":
+            self.lr_scheduler = None
         elif backend_name == "jax":
             self.opt_state = None
 
@@ -251,7 +253,9 @@ def outputs(training, inputs):
         def outputs_losses(training, inputs, targets, losses_fn):
             self.net.train(mode=training)
             if isinstance(inputs, tuple):
-                inputs = tuple(map(lambda x: torch.as_tensor(x).requires_grad_(), inputs))
+                inputs = tuple(
+                    map(lambda x: torch.as_tensor(x).requires_grad_(), inputs)
+                )
             else:
                 inputs = torch.as_tensor(inputs)
                 inputs.requires_grad_()
@@ -283,7 +287,7 @@ def outputs_losses_test(inputs, targets):
         trainable_variables = (
             list(self.net.parameters()) + self.external_trainable_variables
         )
-        self.opt = optimizers.get(
+        self.opt, self.lr_scheduler = optimizers.get(
             trainable_variables, self.opt_name, learning_rate=lr, decay=decay
         )
 
@@ -296,6 +300,8 @@ def closure():
                 return total_loss
 
             self.opt.step(closure)
+            if self.lr_scheduler is not None:
+                self.lr_scheduler.step()
 
         # Callables
         self.outputs = outputs

deepxde/optimizers/pytorch/optimizers.py

@@ -11,13 +11,13 @@ def is_external_optimizer(optimizer):
 
 def get(params, optimizer, learning_rate=None, decay=None):
     """Retrieves an Optimizer instance."""
+    # Custom Optimizer
     if isinstance(optimizer, torch.optim.Optimizer):
-        return optimizer
-
-    if optimizer in ["L-BFGS", "L-BFGS-B"]:
+        optim = optimizer
+    elif optimizer in ["L-BFGS", "L-BFGS-B"]:
         if learning_rate is not None or decay is not None:
             print("Warning: learning rate is ignored for {}".format(optimizer))
-        return torch.optim.LBFGS(
+        optim = torch.optim.LBFGS(
             params,
             lr=1,
             max_iter=LBFGS_options["iter_per_step"],
@@ -27,15 +27,33 @@ def get(params, optimizer, learning_rate=None, decay=None):
             history_size=LBFGS_options["maxcor"],
             line_search_fn=None,
         )
-
-    if learning_rate is None:
-        raise ValueError("No learning rate for {}.".format(optimizer))
-
-    if decay is not None:
-        # TODO: learning rate decay
-        raise NotImplementedError(
-            "learning rate decay to be implemented for backend pytorch."
+    else:
+        if learning_rate is None:
+            raise ValueError("No learning rate for {}.".format(optimizer))
+        if optimizer == "sgd":
+            optim = torch.optim.SGD(params, lr=learning_rate)
+        elif optimizer == "rmsprop":
+            optim = torch.optim.RMSprop(params, lr=learning_rate)
+        elif optimizer == "adam":
+            optim = torch.optim.Adam(params, lr=learning_rate)
+        else:
+            raise NotImplementedError(
+                f"{optimizer} to be implemented for backend pytorch."
+            )
+    lr_scheduler = _get_learningrate_scheduler(optim, decay)
+    return optim, lr_scheduler
+
+
+def _get_learningrate_scheduler(optim, decay):
+    if decay is None:
+        return None
+
+    if decay[0] == "step":
+        return torch.optim.lr_scheduler.StepLR(
+            optim, step_size=decay[1], gamma=decay[2]
         )
-    if optimizer == "adam":
-        return torch.optim.Adam(params, lr=learning_rate)
-    raise NotImplementedError(f"{optimizer} to be implemented for backend pytorch.")
+
+    # TODO: More learning rate scheduler
+    raise NotImplementedError(
+        f"{decay[0]} learning rate scheduler to be implemented for backend pytorch."
+    )