trainer: Add learning rate decay upon plateau

Author: lisavader

trainer/train_model.py

@@ -79,6 +79,12 @@ def get_args():
         default=0.0005,
         help='Minimum improvement in validation metric to reset patience (default: %(default)s)',
     )
+    parser.add_argument(
+        '--learning-rate-decay',
+        type=float,
+        default=1/3,
+        help='Factor to decay learning rate by on plateau (default: %(default)s)',
+    )
     parser.add_argument(
         '--chunks-per-batch',
         type=int,
@@ -279,6 +285,7 @@ def categorical_crossentropy_2d_gene_ml(y_true, y_pred):
     best_val_score = -1.0
     best_epoch = 0
     patience_counter = 0
+    decay_counter = 0
 
     # Evaluation batch size can be smaller than training to reduce peak GPU mem
     EVAL_BATCH_SIZE = max(1, min(BATCH_SIZE, args.eval_batch_size * N_GPUS))
@@ -423,11 +430,20 @@ def print_performance_metrics(indices, max_eval):
 
             # Early stopping: halt if patience exceeded
             if args.early_stopping_patience > 0 and patience_counter >= args.early_stopping_patience:
-                tee(f"\n\033[93mEarly stopping triggered after {patience_counter} evaluations without improvement.\033[0m")
-                tee(f"\033[92mBest epoch: {best_epoch} with validation score: {best_val_score:.4f}\033[0m")
-                tee(f"\033[92mRecommended checkpoint: GeneML{args.context_length}_c{args.dataset_name}_ep{best_epoch}.keras\033[0m")
-                h5f.close()
-                sys.exit(0)
+                # allow learning rate to decay up to 2 times before stopping
+                if decay_counter < 2:
+                    old_lr = K.get_value(model.optimizer.learning_rate)
+                    new_lr = old_lr * args.learning_rate_decay
+                    K.set_value(model.optimizer.learning_rate, new_lr)
+                    tee(f"\n\033[93mLearning rate decayed from {old_lr:.5f} to {new_lr:.5f} after {patience_counter} evaluations without improvement.\033[0m")
+                    patience_counter = 0
+                    decay_counter += 1
+                else:
+                    tee(f"\n\033[93mEarly stopping triggered after {patience_counter} evaluations without improvement.\033[0m")
+                    tee(f"\033[92mBest epoch: {best_epoch} with validation score: {best_val_score:.4f}\033[0m")
+                    tee(f"\033[92mRecommended checkpoint: GeneML{args.context_length}_c{args.dataset_name}_ep{best_epoch}.keras\033[0m")
+                    h5f.close()
+                    sys.exit(0)
         else:
             tee(f"Skipping evaluation this epoch (eval_every={args.eval_every})")
             tee("--- %s seconds ---" % (time.time() - start_time))