From c7133a8bb127fd3773ec9df2fe180b14625d5fef Mon Sep 17 00:00:00 2001
From: Cian Hughes <chughes000@gmail.com>
Date: Fri, 7 Jun 2024 10:58:25 +0100
Subject: [PATCH] Switched loss in expt2 to smooth_l1

Switched from mse_loss to smooth_l1_loss to avoid exploding gradient and
NaNs when using mse_loss.
---
 symbolic_nn_tests/experiment2/__init__.py      | 14 +++++++-------
 symbolic_nn_tests/experiment2/model.py         | 10 +++++-----
 symbolic_nn_tests/experiment2/semantic_loss.py |  2 +-
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/symbolic_nn_tests/experiment2/__init__.py b/symbolic_nn_tests/experiment2/__init__.py
index 1e756e0..4e99af8 100644
--- a/symbolic_nn_tests/experiment2/__init__.py
+++ b/symbolic_nn_tests/experiment2/__init__.py
@@ -40,21 +40,21 @@ def test(train_loss, val_loss, test_loss, version, tensorboard=True, wandb=True)
 
 
 def run(tensorboard: bool = True, wandb: bool = True):
-    from .model import unpacking_mse_loss
+    from .model import unpacking_smooth_l1_loss
     from . import semantic_loss
 
     test(
-        train_loss=unpacking_mse_loss,
-        val_loss=unpacking_mse_loss,
-        test_loss=unpacking_mse_loss,
-        version="mse_loss",
+        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,
     )
     test(
         train_loss=semantic_loss.positive_slope_linear_loss,
-        val_loss=unpacking_mse_loss,
-        test_loss=unpacking_mse_loss,
+        val_loss=unpacking_smooth_l1_loss,
+        test_loss=unpacking_smooth_l1_loss,
         version="positive_slope_linear_loss",
         tensorboard=tensorboard,
         wandb=wandb,
diff --git a/symbolic_nn_tests/experiment2/model.py b/symbolic_nn_tests/experiment2/model.py
index d1d207a..8bb70bd 100644
--- a/symbolic_nn_tests/experiment2/model.py
+++ b/symbolic_nn_tests/experiment2/model.py
@@ -61,15 +61,15 @@ def get_singleton_dataset():
     )
 
 
-def unpacking_mse_loss(out, y):
+def unpacking_smooth_l1_loss(out, y):
     _, y_pred = out
-    return nn.functional.mse_loss(y_pred, y)
+    return nn.functional.smooth_l1_loss(y_pred, y)
 
 
 def main(
-    train_loss=unpacking_mse_loss,
-    val_loss=unpacking_mse_loss,
-    test_loss=unpacking_mse_loss,
+    train_loss=unpacking_smooth_l1_loss,
+    val_loss=unpacking_smooth_l1_loss,
+    test_loss=unpacking_smooth_l1_loss,
     logger=None,
     **kwargs,
 ):
diff --git a/symbolic_nn_tests/experiment2/semantic_loss.py b/symbolic_nn_tests/experiment2/semantic_loss.py
index 0294990..e8d8ec6 100644
--- a/symbolic_nn_tests/experiment2/semantic_loss.py
+++ b/symbolic_nn_tests/experiment2/semantic_loss.py
@@ -66,4 +66,4 @@ def positive_slope_linear_loss(out, y):
     slope_penalty = (torch.nn.functional.softmax(-m * 500.0) + 1).mean()
 
     # Finally, let's get a smooth L1 loss and scale it based on these penalty functions
-    return nn.functional.mse_loss(y_pred, y) * residual_penalty * slope_penalty
+    return nn.functional.smooth_l1_loss(y_pred, y) * residual_penalty * slope_penalty