From a113de0d0ba103b6fb1bfdec69b550147a2a262f Mon Sep 17 00:00:00 2001
From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= <francois@fleuret.org>
Date: Tue, 28 Mar 2023 22:17:16 +0200
Subject: [PATCH] Update

---
 beaver.py | 56 ++++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 47 insertions(+), 9 deletions(-)

diff --git a/beaver.py b/beaver.py
index f395d22..f5b3563 100755
--- a/beaver.py
+++ b/beaver.py
@@ -264,7 +264,7 @@ def oneshot(gpt, learning_rate_scheduler, task):
     learning_rate_scheduler.reset()
 
     for n_epoch in range(args.nb_epochs):
-        learning_rate = learning_rate_scheduler.learning_rate()
+        learning_rate = learning_rate_scheduler.get_learning_rate()
         optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
 
         acc_train_loss, nb_train_samples = 0, 0
@@ -342,7 +342,7 @@ def oneshot(gpt, learning_rate_scheduler, task):
 
 
 class LearningRateScheduler:
-    def learning_rate(self):
+    def get_learning_rate(self):
         pass
 
     def update(self, nb_finished_epochs, loss):
@@ -355,7 +355,8 @@ class LearningRateScheduler:
         return vars(self)
 
     def set_state(self, state):
-        for k, v in state.item():
+        print(f"{state=}")
+        for k, v in state.items():
             setattr(self, k, v)
 
 
@@ -364,12 +365,47 @@ class StepWiseScheduler(LearningRateScheduler):
         self.nb_finished_epochs = 0
         self.schedule = schedule
 
-    def learning_rate(self):
+    def get_learning_rate(self):
         return self.schedule[self.nb_finished_epochs]
 
+    def update(self, nb_finished_epochs, loss):
+        self.nb_finished_epochs = nb_finished_epochs
+
     def reset(self):
         self.nb_finished_epochs = 0
 
+    def get_state(self):
+        return {"nb_finished_epochs": self.nb_finished_epochs}
+
+
+class AutoScheduler(LearningRateScheduler):
+    def __init__(self, learning_rate_init, growth=1.0, degrowth=0.2):
+        self.learning_rate_init = learning_rate_init
+        self.learning_rate = learning_rate_init
+        self.growth = growth
+        self.degrowth = degrowth
+        self.pred_loss = None
+
+    def get_learning_rate(self):
+        return self.learning_rate
+
+    def update(self, nb_finished_epochs, loss):
+        if self.pred_loss is not None:
+            if loss >= self.pred_loss:
+                self.learning_rate *= self.degrowth
+            else:
+                self.learning_rate *= self.growth
+        self.pred_loss = loss
+
+    def reset(self):
+        self.learning_rate = self.learning_rate_init
+
+    def get_state(self):
+        return {
+            "learning_rate_init": self.learning_rate_init,
+            "pred_loss": self.pred_loss,
+        }
+
 
 ######################################################################
 
@@ -589,7 +625,7 @@ log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)")
 ######################################################################
 
 if args.learning_rate_schedule == "auto":
-    pass
+    learning_rate_scheduler = AutoScheduler(args.learning_rate)
 
 elif args.learning_rate_schedule == "cos":
     schedule = {}
@@ -629,6 +665,7 @@ else:
         checkpoint = torch.load(checkpoint_name)
         nb_epochs_finished = checkpoint["nb_epochs_finished"]
         model.load_state_dict(checkpoint["model_state"])
+        learning_rate_scheduler.set_state(checkpoint["learning_rate_scheduler_state"])
         torch.set_rng_state(checkpoint["rng_state"])
         if torch.cuda.is_available():
             torch.cuda.set_rng_state(checkpoint["cuda_rng_state"])
@@ -638,9 +675,9 @@ else:
     except FileNotFoundError:
         log_string("starting from scratch.")
 
-    except:
-        log_string("error when loading the checkpoint.")
-        exit(1)
+    # except:
+    # log_string("error when loading the checkpoint.")
+    # exit(1)
 
 ######################################################################
 
@@ -673,7 +710,7 @@ if nb_epochs_finished >= args.nb_epochs:
 learning_rate_scheduler.reset()
 
 for n_epoch in range(nb_epochs_finished, args.nb_epochs):
-    learning_rate = learning_rate_scheduler.learning_rate()
+    learning_rate = learning_rate_scheduler.get_learning_rate()
 
     log_string(f"learning_rate {learning_rate}")
 
@@ -721,6 +758,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
     checkpoint = {
         "nb_epochs_finished": n_epoch + 1,
         "model_state": model.state_dict(),
+        "learning_rate_scheduler_state": learning_rate_scheduler.get_state(),
         "rng_state": torch.get_rng_state(),
     }
 
-- 
2.39.5