add huber loss

MannLabs · May 29, 2024 · f2b2e2d · f2b2e2d
1 parent f2139e1
commit f2b2e2d
Showing 1 changed file with 3 additions and 7 deletions.
diff --git a/src/sparcscore/ml/plmodels.py b/src/sparcscore/ml/plmodels.py
@@ -206,9 +206,7 @@ def training_step(self, batch):
         data, target = batch
         target = target.unsqueeze(1)
         output = self.network(data) # Forward pass, only one output
-        loss = F.mse_loss(output, target) # L2 loss
-
-        # accuracy metrics for regression???
+        loss = F.huber_loss(output, target, delta=1.0, reduction='mean') # consider looking at parameters again
 
         self.log('loss/train', loss, on_step=False, on_epoch=True, prog_bar=True)
         self.log('mse/train', self.mse(output, target), on_epoch=True, prog_bar=True)
@@ -220,9 +218,7 @@ def validation_step(self, batch):
         data, target = batch
         target = target.unsqueeze(1)
         output = self.network(data)
-        loss = F.mse_loss(output, target)
-
-         # accuracy metrics for regression???
+        loss = F.huber_loss(output, target, delta=1.0, reduction='mean')
 
         self.log('loss/val', loss, on_step=False, on_epoch=True, prog_bar=True)
         self.log('mse/val', self.mse(output, target), on_epoch=True, prog_bar=True)
@@ -234,7 +230,7 @@ def test_step(self, batch):
         data, target = batch
         target = target.unsqueeze(1)
         output = self.network(data)
-        loss = F.mse_loss(output, target)
+        loss = F.huber_loss(output, target, delta=1.0, reduction='mean')
 
          # accuracy metrics for regression???