exclude main loss from monitored loss types

equitrain/argparser.py

@@ -384,6 +384,16 @@ def check_args_complete(args: argparse.ArgumentParser, script_type: str):
             raise ValueError(f'Unexpected arguments: {extra}')
 
 
+def get_loss_monitor(args: argparse.ArgumentParser) -> list[str]:
+    # Create list of loss types
+    loss_monitor = [item.strip().lower() for item in args['loss_monitor'].split(',')]
+
+    if args.loss_type in loss_monitor:
+        loss_monitor.remove(args.loss_type)
+
+    return loss_monitor
+
+
 class ArgumentError(ValueError):
     """Custom exception raised when invalid or missing argument is present."""
 

equitrain/loss_fn.py

@@ -1,5 +1,6 @@
 import torch
 
+from equitrain.argparser import get_loss_monitor
 from equitrain.data.scatter import scatter_mean
 from equitrain.loss import Loss, LossCollection
 
@@ -228,7 +229,7 @@ def __init__(self, **args):
 
         # Additional loss metrics
         self.loss_fns = {}
-        for loss_type in args['loss_monitor'].split(','):
+        for loss_type in get_loss_monitor(args):
             args_new = {**args, 'loss_type': loss_type}
             self.loss_fns[loss_type] = LossFn(**args_new)
 

equitrain/loss_metrics.py

@@ -1,3 +1,4 @@
+from equitrain.argparser import get_loss_monitor
 from equitrain.loss import LossCollection
 
 
@@ -141,13 +142,13 @@ def update(self, loss, epoch):
 class LossMetrics(dict):
     def __init__(self, args):
         self.main = LossMetric(args)
-        self.main_type = args.loss_type
-        for loss_type in args.loss_monitor.split(','):
+        self.main_type = args.loss_type.lower()
+        for loss_type in get_loss_monitor(args):
             self[loss_type] = LossMetric(args)
 
     def update(self, loss: LossCollection):
         self.main.update(loss.main)
-        for loss_type, loss_metric in self.items():
+        for loss_type, _ in self.items():
             self[loss_type].update(loss[loss_type])
 
     def log(self, logger, mode: str, epoch=None, step=None, time=None, lr=None):