optional monitoring of validation loss for the scheduler

equitrain/argparser.py

@@ -170,6 +170,7 @@ def add_optimizer_args(parser: argparse.ArgumentParser) -> argparse.ArgumentPars
     )
     parser.add_argument(
         '--plateau-mode',
+        choices=['min', 'max'],
         type=str,
         default='min',
         help='One of min, max. In min mode, lr will be reduced when the quantity monitored has stopped decreasing; in max mode it will be reduced when the quantity monitored has stopped increasing (default: min)',
@@ -276,9 +277,17 @@ def get_args_parser(script_type: str) -> argparse.ArgumentParser:
         parser.add_argument(
             '--scheduler', help='LR scheduler type', type=str, default='plateau'
         )
+        parser.add_argument(
+            '--scheduler_monitor',
+            help='Loss monitored by the scheduler [train (default), val]',
+            choices=['train', 'val'],
+            type=str,
+            default='train',
+        )
         parser.add_argument(
             '--loss-type',
             help='Type of loss function [mae, smooth-l1, mse, huber (default)]',
+            choices=['mae', 'smooth-l1', 'mse', 'huber'],
             type=str,
             default='huber',
         )

equitrain/train.py

@@ -398,7 +398,10 @@ def _train_with_accelerator(args, accelerator: Accelerator):
                 )
 
         if lr_scheduler is not None:
-            lr_scheduler.step(metric=train_loss.main['total'].avg, epoch=epoch)
+            if args.scheduler_monitor == 'train':
+                lr_scheduler.step(metric=train_loss.main['total'].avg, epoch=epoch)
+            if args.scheduler_monitor == 'val':
+                lr_scheduler.step(metric=val_loss.main['total'].avg, epoch=epoch)
 
             if last_lr is not None and last_lr != lr_scheduler.get_last_lr()[0]:
                 logger.log(