Backend PyTorch: Add L1 regularizers

deepxde/model.py

@@ -278,6 +278,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():
@@ -313,6 +324,11 @@ def outputs_losses(training, inputs, targets, auxiliary_vars, losses_fn):
             losses = losses_fn(targets, outputs_, loss_fn, inputs, self, aux=aux)
             if not isinstance(losses, list):
                 losses = [losses]
+            if l1_factor > 0:
+                l1_loss = torch.sum(
+                    torch.stack([torch.sum(p.abs()) for p in self.net.parameters()])
+                )
+                losses.append(l1_factor * l1_loss)
             losses = torch.stack(losses)
             # Weighted losses
             if self.loss_weights is not None:
@@ -331,23 +347,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."
@@ -363,7 +370,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):