Added trainable residual penalty & logging for it

symbolic_nn_tests/experiment2/__init__.py

@@ -20,10 +20,14 @@ def test(train_loss, val_loss, test_loss, version, tensorboard=True, wandb=True)
         import wandb as _wandb
         from lightning.pytorch.loggers import WandbLogger
 
+        if isinstance(wandb, WandbLogger):
+            wandb_logger = wandb
+        else:
             wandb_logger = WandbLogger(
                 project="Symbolic_NN_Tests",
                 name=version,
                 dir="wandb",
+                log_model="all",
             )
         logger.append(wandb_logger)
 
@@ -43,19 +47,33 @@ def run(tensorboard: bool = True, wandb: bool = True):
     from .model import unpacking_smooth_l1_loss
     from . import semantic_loss
 
+    # test(
+    #     train_loss=unpacking_smooth_l1_loss,
+    #     val_loss=unpacking_smooth_l1_loss,
+    #     test_loss=unpacking_smooth_l1_loss,
+    #     version="smooth_l1_loss",
+    #     tensorboard=tensorboard,
+    #     wandb=wandb,
+    # )
+
+    version = "positive_slope_linear_loss"
+    if wandb:
+        from lightning.pytorch.loggers import WandbLogger
+
+        wandb_logger = WandbLogger(
+            project="Symbolic_NN_Tests",
+            name=version,
+            dir="wandb",
+            log_model="all",
+        )
+    else:
+        wandb_logger = wandb
+
     test(
-        train_loss=unpacking_smooth_l1_loss,
+        train_loss=semantic_loss.positive_slope_linear_loss(wandb_logger, version),
         val_loss=unpacking_smooth_l1_loss,
         test_loss=unpacking_smooth_l1_loss,
-        version="smooth_l1_loss",
+        version=version,
         tensorboard=tensorboard,
-        wandb=wandb,
-    )
-    test(
-        train_loss=semantic_loss.positive_slope_linear_loss,
-        val_loss=unpacking_smooth_l1_loss,
-        test_loss=unpacking_smooth_l1_loss,
-        version="positive_slope_linear_loss",
-        tensorboard=tensorboard,
-        wandb=wandb,
+        wandb=wandb_logger,
     )

symbolic_nn_tests/experiment2/model.py

@@ -49,7 +49,7 @@ def get_singleton_dataset():
     from symbolic_nn_tests.experiment2.dataset import collate, pubchem
 
     return create_dataset(
-        dataset=pubchem, collate_fn=collate, batch_size=512, shuffle=True
+        dataset=pubchem, collate_fn=collate, batch_size=256, shuffle=True
     )
 
 

symbolic_nn_tests/experiment2/semantic_loss.py

@@ -18,7 +18,10 @@ import torch
 # proportionality.
 
 
-def positive_slope_linear_loss(out, y):
+def positive_slope_linear_loss(wandb_logger=None, version="", device="cuda"):
+    a = nn.Parameter(data=torch.randn(1), requires_grad=True).to(device)
+
+    def f(out, y):
         x, y_pred = out
         x0, x1 = x
 
@@ -50,11 +53,17 @@ def positive_slope_linear_loss(out, y):
             .mean()
         )
 
-    # We also need to calculate a penalty that incentivizes a positive slope. For this, im using softplus
+        # We also need to calculate a penalty that incentivizes a positive slope. For this, im using relu
         # to scale the slope as it will penalise negative slopes without just creating a reward hack for
         # maximizing slope.
-    slope_penalty = (nn.functional.softplus(-m) + 1).mean()
+        slope_penalty = (nn.functional.relu(a * (-m)) + 1).mean()
+
+        if wandb_logger:
+            wandb_logger.log_metrics({f"{version}-a": a})
 
         # Finally, let's get a smooth L1 loss and scale it based on these penalty functions
-    return nn.functional.smooth_l1_loss(y_pred, y) * residual_penalty * slope_penalty
+        return (
+            nn.functional.smooth_l1_loss(y_pred, y) * residual_penalty * slope_penalty
+        )
 
+    return f