give diffusion prior trainer cosine annealing lr too

Author: lucidrains

dalle2_pytorch/trainer.py

@@ -182,6 +182,7 @@ def __init__(
         max_grad_norm = None,
         group_wd_params = True,
         warmup_steps = 1,
+        cosine_decay_max_steps = None,
         **kwargs
     ):
         super().__init__()
@@ -233,8 +234,11 @@ def __init__(
             **self.optim_kwargs,
             **kwargs
         )
-        
-        self.scheduler = LambdaLR(self.optimizer, lr_lambda = lambda _: 1.0)
+
+        if exists(cosine_decay_max_steps):
+            self.scheduler = CosineAnnealingLR(optimizer, T_max = cosine_decay_max_steps)
+        else:
+            self.scheduler = LambdaLR(self.optimizer, lr_lambda = lambda _: 1.0)
 
         self.warmup_scheduler = warmup.LinearWarmup(self.optimizer, warmup_period = warmup_steps) if exists(warmup_steps) else None
 
@@ -271,6 +275,7 @@ def save(self, path, overwrite = True, **kwargs):
             # FIXME: LambdaLR can't be saved due to pickling issues
             save_obj = dict(
                 optimizer = self.optimizer.state_dict(),
+                scheduler = self.scheduler.state_dict(),
                 warmup_scheduler = self.warmup_scheduler,
                 model = self.accelerator.unwrap_model(self.diffusion_prior).state_dict(),
                 version = version.parse(__version__),
@@ -317,7 +322,9 @@ def load(self, path_or_state, overwrite_lr = True, strict = True):
         # unwrap the model when loading from checkpoint
         self.accelerator.unwrap_model(self.diffusion_prior).load_state_dict(loaded_obj['model'], strict = strict)
         self.step.copy_(torch.ones_like(self.step, device=self.device) * loaded_obj['step'].to(self.device))
+
         self.optimizer.load_state_dict(loaded_obj['optimizer'])
+        self.scheduler.load_state_dict(loaded_obj['scheduler'])
 
         # set warmupstep
         if exists(self.warmup_scheduler):

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.8.0'
+__version__ = '1.8.1'