Backend PyTorch supports Adam optimizer

Author: lululxvi

deepxde/model.py

@@ -189,7 +189,9 @@ def outputs_losses(data_id, inputs, targets):
                 losses = compute_losses(targets, outputs)
                 return outputs, losses
 
-            opt = torch.optim.Adam(self.net.parameters(), lr=lr)
+            opt = optimizers.get(
+                self.net.parameters(), self.opt_name, learning_rate=lr, decay=decay
+            )
 
             def train_step(data_id, inputs, targets):
                 _, losses = outputs_losses(data_id, inputs, targets)
@@ -404,6 +406,7 @@ def predict(self, x, operator=None, callbacks=None):
         self.callbacks.set_model(self)
         self.callbacks.on_predict_begin()
         # TODO: use self._run for tensorflow
+        # TODO: predict operator with auxiliary_vars
         if backend_name == "tensorflow.compat.v1":
             if operator is None:
                 op = self.net.outputs

deepxde/optimizers/pytorch/optimizers.py

@@ -1,5 +1,21 @@
-__all__ = ["is_external_optimizer"]
+import torch
+
+
+__all__ = ["get", "is_external_optimizer"]
 
 
 def is_external_optimizer(optimizer):
     return False
+
+
+def get(params, optimizer, learning_rate=None, decay=None):
+    """Retrieves an Optimizer instance."""
+    if isinstance(optimizer, torch.optim.Optimizer):
+        return optimizer
+    if learning_rate is None:
+        raise ValueError("No learning rate for {}.".format(optimizer))
+
+    # TODO: decay
+    if optimizer == "adam":
+        return torch.optim.Adam(params, lr=learning_rate)
+    raise NotImplementedError(f"{optimizer} to be implemented for backend pytorch.")