Make NN training robust to NaN (#237)

Author: RemiLehe

ml/Neural_Net_Classes.py

@@ -23,6 +23,35 @@ def __call__(self, val_loss):
                 self.early_stop = True
 
 
+def nan_mse_loss( target, pred ):
+    """
+    Custom MSE loss that ignores NaN values in targets
+    (Here, NaN often correspond to missing values in the target data)
+
+    Args:
+        target: target values (may contain NaN)
+        pred: predicted values
+
+    Returns:
+        mean squared error ignoring NaN values
+    """
+    # Compute squared differences
+    squared_diff = (pred - target) ** 2
+
+    # Use nanmean to ignore NaN values
+    mse_loss = torch.nanmean(squared_diff)
+
+    # Prevent NaN from contaminating backpropagation
+    # See https://github.com/pytorch/pytorch/issues/4132
+    if pred.requires_grad:
+        nan_mask = torch.isnan(squared_diff)
+        def mask_grad_hook(grad):
+            return torch.where(nan_mask, 0, grad)
+        pred.register_hook(mask_grad_hook)
+
+    return mse_loss
+
+
 class CombinedNN(nn.Module):
     """
     Model that trains a 5 layer neural network and a calibration layer
@@ -52,7 +81,7 @@ def __init__(self, input_size, output_size, hidden_size=20,
         self.sim_to_exp_calibration_weight = nn.Parameter(torch.ones(output_size))
         self.sim_to_exp_calibration_bias = nn.Parameter(torch.zeros(output_size))
 
-        self.criterion = nn.MSELoss()
+        # Use custom loss function instead of nn.MSELoss()
         self.optimizer = optim.Adam(self.parameters(), lr=learning_rate)
         self.scheduler = ReduceLROnPlateau(self.optimizer, 'min',
                                            factor=factor, patience=patience_LRreduction, threshold=threshold)
@@ -91,10 +120,10 @@ def train_model(self, sim_inputs, sim_targets,
             loss = 0
             if len(sim_inputs) > 0:
                 sim_outputs = self(sim_inputs)
-                loss += self.criterion( sim_targets, sim_outputs )
+                loss += nan_mse_loss( sim_targets, sim_outputs )
             if len(exp_inputs) > 0:
                 exp_outputs = self.calibrate( self(exp_inputs) )
-                loss += self.criterion( exp_targets, exp_outputs )
+                loss += nan_mse_loss( exp_targets, exp_outputs )
             loss.backward()
 
             self.optimizer.step()
@@ -107,10 +136,10 @@ def train_model(self, sim_inputs, sim_targets,
                 val_loss = 0
                 if len(sim_inputs_val) > 0:
                     sim_outputs_val = self(sim_inputs_val)
-                    val_loss += self.criterion( sim_targets_val, sim_outputs_val )
+                    val_loss += nan_mse_loss( sim_targets_val, sim_outputs_val )
                 if len(exp_inputs_val) > 0:
                     exp_outputs_val = self(exp_inputs_val)
-                    val_loss += self.criterion( exp_targets_val, exp_outputs_val )
+                    val_loss += nan_mse_loss( exp_targets_val, exp_outputs_val )
 
 
             if(epoch+1) % (num_epochs/10) == 0: