Backend PyTorch: Add L1 and L1+L2 regularizers (#1905)

Author: vl-dud

deepxde/model.py

@@ -279,6 +279,17 @@ def build_loss():
     def _compile_pytorch(self, lr, loss_fn, decay):
         """pytorch"""
 
+        l1_factor, l2_factor = 0, 0
+        if self.net.regularizer is not None:
+            if self.net.regularizer[0] == "l1":
+                l1_factor = self.net.regularizer[1]
+            elif self.net.regularizer[0] == "l2":
+                l2_factor = self.net.regularizer[1]
+            else:
+                raise NotImplementedError(
+                    f"{self.net.regularizer[0]} regularizer hasn't been implemented for backend pytorch."
+                )
+
         def outputs(training, inputs):
             self.net.train(mode=training)
             with torch.no_grad():
@@ -318,6 +329,11 @@ def outputs_losses(training, inputs, targets, auxiliary_vars, losses_fn):
             # Weighted losses
             if self.loss_weights is not None:
                 losses *= torch.as_tensor(self.loss_weights)
+            if l1_factor > 0:
+                l1_loss = l1_factor * torch.sum(
+                    torch.stack([torch.sum(p.abs()) for p in self.net.parameters()])
+                )
+                losses = torch.cat([losses, l1_loss.unsqueeze(0)])
             # Clear cached Jacobians and Hessians.
             grad.clear()
             return outputs_, losses
@@ -332,23 +348,14 @@ def outputs_losses_test(inputs, targets, auxiliary_vars):
                 False, inputs, targets, auxiliary_vars, self.data.losses_test
             )
 
-        weight_decay = 0
-        if self.net.regularizer is not None:
-            if self.net.regularizer[0] != "l2":
-                raise NotImplementedError(
-                    f"{self.net.regularizer[0]} regularization to be implemented for "
-                    "backend pytorch"
-                )
-            weight_decay = self.net.regularizer[1]
-
         optimizer_params = self.net.parameters()
         if self.external_trainable_variables:
             # L-BFGS doesn't support per-parameter options.
             if self.opt_name in ["L-BFGS", "L-BFGS-B"]:
                 optimizer_params = (
                     list(optimizer_params) + self.external_trainable_variables
                 )
-                if weight_decay > 0:
+                if l2_factor > 0:
                     print(
                         "Warning: L2 regularization will also be applied to external_trainable_variables. "
                         "Ensure this is intended behavior."
@@ -364,7 +371,7 @@ def outputs_losses_test(inputs, targets, auxiliary_vars):
             self.opt_name,
             learning_rate=lr,
             decay=decay,
-            weight_decay=weight_decay,
+            weight_decay=l2_factor,
         )
 
         def train_step(inputs, targets, auxiliary_vars):