Add option to change learning rate scheduler and made it easier to ad…

…d a new one.

casanovo/config.py

@@ -53,6 +53,7 @@ class Config:
         residues=dict,
         n_log=int,
         tb_summarywriter=str,
+        lr_schedule=str,
         warmup_iters=int,
         max_iters=int,
         learning_rate=float,

casanovo/config.yaml

@@ -84,6 +84,8 @@ dim_intensity:
 custom_encoder:
 # Max decoded peptide length
 max_length: 100
+# Type of learning rate schedule to use. One of {constant, linear, cosine}.
+lr_schedule: "cosine"
 # Number of warmup iterations for learning rate scheduler
 warmup_iters: 100_000
 # Max number of iterations for learning rate scheduler

casanovo/denovo/model.py

@@ -110,6 +110,7 @@ def __init__(
         tb_summarywriter: Optional[
             torch.utils.tensorboard.SummaryWriter
         ] = None,
+        lr_schedule = None,
         warmup_iters: int = 100_000,
         max_iters: int = 600_000,
         out_writer: Optional[ms_io.MztabWriter] = None,
@@ -143,6 +144,7 @@ def __init__(
         self.softmax = torch.nn.Softmax(2)
         self.celoss = torch.nn.CrossEntropyLoss(ignore_index=0)
         # Optimizer settings.
+        self.lr_schedule = lr_schedule
         self.warmup_iters = warmup_iters
         self.max_iters = max_iters
         self.opt_kwargs = kwargs
@@ -964,13 +966,35 @@ def configure_optimizers(
         """
         optimizer = torch.optim.Adam(self.parameters(), **self.opt_kwargs)
         # Apply learning rate scheduler per step.
-        lr_scheduler = CosineWarmupScheduler(
-            optimizer, warmup=self.warmup_iters, max_iters=self.max_iters
-        )
+        lr_schedulers = [
+            torch.optim.lr_scheduler.LinearLR(
+                optimizer, 
+                start_factor=1e-10, 
+                total_iters=self.warmup_iters)
+            ]
+        if self.lr_schedule == 'cosine':
+            lr_schedulers.append(
+                CosineScheduler(
+                    optimizer, 
+                    warmup=self.warmup_iters, 
+                    max_iters=self.max_iters
+                    )
+            )
+        elif self.lr_schedule == 'linear':
+            lr_schedulers.append(
+                torch.optim.lr_scheduler.LinearLR(
+                    optimizer, 
+                    start_factor=1, 
+                    end_factor=0,
+                    total_iters=self.max_iters
+                    )
+            )
+
+        lr_scheduler = torch.optim.lr_scheduler.ChainedScheduler(lr_schedulers)
         return [optimizer], {"scheduler": lr_scheduler, "interval": "step"}
 
 
-class CosineWarmupScheduler(torch.optim.lr_scheduler._LRScheduler):
+class CosineScheduler(torch.optim.lr_scheduler._LRScheduler):
     """
     Learning rate scheduler with linear warm up followed by cosine shaped decay.
 
@@ -985,9 +1009,9 @@ class CosineWarmupScheduler(torch.optim.lr_scheduler._LRScheduler):
     """
 
     def __init__(
-        self, optimizer: torch.optim.Optimizer, warmup: int, max_iters: int
+        self, optimizer: torch.optim.Optimizer, max_iters: int
     ):
-        self.warmup, self.max_iters = warmup, max_iters
+        self.max_iters = max_iters
         super().__init__(optimizer)
 
     def get_lr(self):
@@ -996,8 +1020,6 @@ def get_lr(self):
 
     def get_lr_factor(self, epoch):
         lr_factor = 0.5 * (1 + np.cos(np.pi * epoch / self.max_iters))
-        if epoch <= self.warmup:
-            lr_factor *= epoch / self.warmup
         return lr_factor
 
 

casanovo/denovo/model_runner.py

@@ -222,6 +222,7 @@ def initialize_model(self, train: bool) -> None:
             top_match=self.config.top_match,
             n_log=self.config.n_log,
             tb_summarywriter=self.config.tb_summarywriter,
+            lr_schedule=self.config.tb_summarywriter,
             warmup_iters=self.config.warmup_iters,
             max_iters=self.config.max_iters,
             lr=self.config.learning_rate,