Rename max_iters to cosine_schedule_period_iters

casanovo/config.py

@@ -2,6 +2,7 @@
 
 import logging
 import shutil
+import warnings
 from pathlib import Path
 from typing import Optional, Dict, Callable, Tuple, Union
 
@@ -33,6 +34,11 @@ class Config:
     """
 
     _default_config = Path(__file__).parent / "config.yaml"
+    # FIXME: This contains deprecated config options to be removed in the next
+    #  major version update.
+    _config_deprecated_remap = dict(
+        max_iters="cosine_schedule_period_iters",
+    )
     _config_types = dict(
         random_seed=int,
         n_peaks=int,
@@ -56,7 +62,7 @@ class Config:
         tb_summarywriter=str,
         train_label_smoothing=float,
         warmup_iters=int,
-        max_iters=int,
+        cosine_schedule_period_iters=int,
         learning_rate=float,
         weight_decay=float,
         train_batch_size=int,
@@ -84,6 +90,15 @@ def __init__(self, config_file: Optional[str] = None):
         else:
             with Path(config_file).open() as f_in:
                 self._user_config = yaml.safe_load(f_in)
+                # Remap deprecated config entries.
+                for old, new in self._config_deprecated_remap.items():
+                    if old in self._user_config:
+                        self._user_config[new] = self._user_config.pop(old)
+                        warnings.warn(
+                            f"Deprecated config option '{old}' remapped to "
+                            f"'{new}'",
+                            DeprecationWarning,
+                        )
                 # Check for missing entries in config file.
                 config_missing = self._params.keys() - self._user_config.keys()
                 if len(config_missing) > 0:

casanovo/config.yaml

@@ -4,103 +4,102 @@
 ###
 
 ###
-# The following parameters can be modified when running inference or
-# when fine-tuning an existing Casanovo model.
+# The following parameters can be modified when running inference or when
+# fine-tuning an existing Casanovo model.
 ###
 
-# Max absolute difference allowed with respect to observed precursor m/z
+# Max absolute difference allowed with respect to observed precursor m/z.
 # Predictions outside the tolerance range are assigned a negative peptide score.
 precursor_mass_tol: 50  # ppm
-# Isotopes to consider when comparing predicted and observed precursor m/z's
+# Isotopes to consider when comparing predicted and observed precursor m/z's.
 isotope_error_range: [0, 1]
-# The minimum length of predicted peptides
+# The minimum length of predicted peptides.
 min_peptide_len: 6
-# Number of spectra in one inference batch
+# Number of spectra in one inference batch.
 predict_batch_size: 1024
-# Number of beams used in beam search
+# Number of beams used in beam search.
 n_beams: 1
-# Number of PSMs for each spectrum
+# Number of PSMs for each spectrum.
 top_match: 1
 # The hardware accelerator to use. Must be one of:
-# "cpu", "gpu", "tpu", "ipu", "hpu", "mps", or "auto"
+# "cpu", "gpu", "tpu", "ipu", "hpu", "mps", or "auto".
 accelerator: "auto"
-# The devices to use. Can be set to a positive number int,
-# or the value -1 to indicate all available devices should be used,
-# If left empty, the appropriate number will be automatically
-# selected for automatic selected on the chosen accelerator.
+# The devices to use. Can be set to a positive number int, or the value -1 to
+# indicate all available devices should be used. If left empty, the appropriate
+# number will be automatically selected for based on the chosen accelerator.
 devices:
 
 ###
 # The following parameters should only be modified if you are training a new
 # Casanovo model from scratch.
 ###
 
-# Random seed to ensure reproducible results
+# Random seed to ensure reproducible results.
 random_seed: 454
 
 # OUTPUT OPTIONS
-# Logging frequency in training steps
+# Logging frequency in training steps.
 n_log: 1
-# Tensorboard directory to use for keeping track of training metrics
+# Tensorboard directory to use for keeping track of training metrics.
 tb_summarywriter:
-# Save the top k model checkpoints during training. -1 saves all, and
-# leaving this field empty saves none.
+# Save the top k model checkpoints during training. -1 saves all, and leaving
+# this field empty saves none.
 save_top_k: 5
-# Path to saved checkpoints
+# Path to saved checkpoints.
 model_save_folder_path: ""
-# Model validation and checkpointing frequency in training steps
+# Model validation and checkpointing frequency in training steps.
 val_check_interval: 50_000
 
 # SPECTRUM PROCESSING OPTIONS
-# Number of the most intense peaks to retain, any remaining peaks are discarded
+# Number of the most intense peaks to retain, any remaining peaks are discarded.
 n_peaks: 150
-# Min peak m/z allowed, peaks with smaller m/z are discarded
+# Min peak m/z allowed, peaks with smaller m/z are discarded.
 min_mz: 50.0
-# Max peak m/z allowed, peaks with larger m/z are discarded
+# Max peak m/z allowed, peaks with larger m/z are discarded.
 max_mz: 2500.0
-# Min peak intensity allowed, less intense peaks are discarded
+# Min peak intensity allowed, less intense peaks are discarded.
 min_intensity: 0.01
-# Max absolute m/z difference allowed when removing the precursor peak
+# Max absolute m/z difference allowed when removing the precursor peak.
 remove_precursor_tol: 2.0  # Da
-# Max precursor charge allowed, spectra with larger charge are skipped
+# Max precursor charge allowed, spectra with larger charge are skipped.
 max_charge: 10
 
 # MODEL ARCHITECTURE OPTIONS
-# Dimensionality of latent representations, i.e. peak embeddings
+# Dimensionality of latent representations, i.e. peak embeddings.
 dim_model: 512
-# Number of attention heads
+# Number of attention heads.
 n_head: 8
-# Dimensionality of fully connected layers
+# Dimensionality of fully connected layers.
 dim_feedforward: 1024
-# Number of transformer layers in spectrum encoder and peptide decoder
+# Number of transformer layers in spectrum encoder and peptide decoder.
 n_layers: 9
-# Dropout rate for model weights
+# Dropout rate for model weights.
 dropout: 0.0
-# Number of dimensions to use for encoding peak intensity
-# Projected up to ``dim_model`` by default and summed with the peak m/z encoding
+# Number of dimensions to use for encoding peak intensity.
+# Projected up to `dim_model` by default and summed with the peak m/z encoding.
 dim_intensity:
-# Max decoded peptide length
+# Max decoded peptide length.
 max_length: 100
-# Number of warmup iterations for learning rate scheduler
+# The number of iterations for the linear warm-up of the learning rate.
 warmup_iters: 100_000
-# Max number of iterations for learning rate scheduler
-max_iters: 600_000
-# Learning rate for weight updates during training
+# The number of iterations for the cosine period of the learning rate.
+cosine_schedule_period_iters: 600_000
+# Learning rate for weight updates during training.
 learning_rate: 5e-4
-# Regularization term for weight updates
+# Regularization term for weight updates.
 weight_decay: 1e-5
-# Amount of label smoothing when computing the training loss
+# Amount of label smoothing when computing the training loss.
 train_label_smoothing: 0.01
 
 # TRAINING/INFERENCE OPTIONS
-# Number of spectra in one training batch
+# Number of spectra in one training batch.
 train_batch_size: 32
-# Max number of training epochs
+# Max number of training epochs.
 max_epochs: 30
-# Number of validation steps to run before training begins
+# Number of validation steps to run before training begins.
 num_sanity_val_steps: 0
-# Calculate peptide and amino acid precision during training. this
-# is expensive, so we recommend against it.
+# Calculate peptide and amino acid precision during training.
+# This is expensive, so we recommend against it.
 calculate_precision: False
 
 # AMINO ACID AND MODIFICATION VOCABULARY

casanovo/denovo/model.py

@@ -46,7 +46,7 @@ class Spec2Pep(pl.LightningModule, ModelMixin):
         linear layer, then summed with the m/z encoding for each peak.
     max_length : int
         The maximum peptide length to decode.
-    residues: Union[Dict[str, float], str]
+    residues : Union[Dict[str, float], str]
         The amino acid dictionary and their masses. By default ("canonical) this
         is only the 20 canonical amino acids, with cysteine carbamidomethylated.
         If "massivekb", this dictionary will include the modifications found in
@@ -65,24 +65,24 @@ class Spec2Pep(pl.LightningModule, ModelMixin):
         < precursor_mass_tol`
     min_peptide_len : int
         The minimum length of predicted peptides.
-    n_beams: int
+    n_beams : int
         Number of beams used during beam search decoding.
-    top_match: int
+    top_match : int
         Number of PSMs to return for each spectrum.
     n_log : int
         The number of epochs to wait between logging messages.
-    tb_summarywriter: Optional[str]
+    tb_summarywriter : Optional[str]
         Folder path to record performance metrics during training. If ``None``,
         don't use a ``SummaryWriter``.
-    train_label_smoothing: float
+    train_label_smoothing : float
         Smoothing factor when calculating the training loss.
-    warmup_iters: int
-        The number of warm up iterations for the learning rate scheduler.
-    max_iters: int
-        The total number of iterations for the learning rate scheduler.
-    out_writer: Optional[str]
+    warmup_iters : int
+        The number of iterations for the linear warm-up of the learning rate.
+    cosine_schedule_period_iters : int
+        The number of iterations for the cosine period of the learning rate.
+    out_writer : Optional[str]
         The output writer for the prediction results.
-    calculate_precision: bool
+    calculate_precision : bool
         Calculate the validation set precision during training.
         This is expensive.
     **kwargs : Dict
@@ -111,7 +111,7 @@ def __init__(
         ] = None,
         train_label_smoothing: float = 0.01,
         warmup_iters: int = 100_000,
-        max_iters: int = 600_000,
+        cosine_schedule_period_iters: int = 600_000,
         out_writer: Optional[ms_io.MztabWriter] = None,
         calculate_precision: bool = False,
         **kwargs: Dict,
@@ -144,7 +144,7 @@ def __init__(
         self.val_celoss = torch.nn.CrossEntropyLoss(ignore_index=0)
         # Optimizer settings.
         self.warmup_iters = warmup_iters
-        self.max_iters = max_iters
+        self.cosine_schedule_period_iters = cosine_schedule_period_iters
         self.opt_kwargs = kwargs
 
         # Data properties.
@@ -960,39 +960,45 @@ 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
+            optimizer, self.warmup_iters, self.cosine_schedule_period_iters
         )
         return [optimizer], {"scheduler": lr_scheduler, "interval": "step"}
 
 
 class CosineWarmupScheduler(torch.optim.lr_scheduler._LRScheduler):
     """
-    Learning rate scheduler with linear warm up followed by cosine shaped decay.
+    Learning rate scheduler with linear warm-up followed by cosine shaped decay.
 
     Parameters
     ----------
     optimizer : torch.optim.Optimizer
         Optimizer object.
-    warmup : int
-        The number of warm up iterations.
-    max_iters : torch.optim
-        The total number of iterations.
+    warmup_iters : int
+        The number of iterations for the linear warm-up of the learning rate.
+    cosine_schedule_period_iters : int
+        The number of iterations for the cosine period of the learning rate.
     """
 
     def __init__(
-        self, optimizer: torch.optim.Optimizer, warmup: int, max_iters: int
+        self,
+        optimizer: torch.optim.Optimizer,
+        warmup_iters: int,
+        cosine_schedule_period_iters: int,
     ):
-        self.warmup, self.max_iters = warmup, max_iters
+        self.warmup_iters = warmup_iters
+        self.cosine_schedule_period_iters = cosine_schedule_period_iters
         super().__init__(optimizer)
 
     def get_lr(self):
         lr_factor = self.get_lr_factor(epoch=self.last_epoch)
         return [base_lr * lr_factor for base_lr in self.base_lrs]
 
     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
+        lr_factor = 0.5 * (
+            1 + np.cos(np.pi * epoch / self.cosine_schedule_period_iters)
+        )
+        if epoch <= self.warmup_iters:
+            lr_factor *= epoch / self.warmup_iters
         return lr_factor
 
 

casanovo/denovo/model_runner.py

@@ -226,7 +226,7 @@ def initialize_model(self, train: bool) -> None:
             tb_summarywriter=self.config.tb_summarywriter,
             train_label_smoothing=self.config.train_label_smoothing,
             warmup_iters=self.config.warmup_iters,
-            max_iters=self.config.max_iters,
+            cosine_schedule_period_iters=self.config.cosine_schedule_period_iters,
             lr=self.config.learning_rate,
             weight_decay=self.config.weight_decay,
             out_writer=self.writer,
@@ -245,7 +245,7 @@ def initialize_model(self, train: bool) -> None:
             tb_summarywriter=self.config.tb_summarywriter,
             train_label_smoothing=self.config.train_label_smoothing,
             warmup_iters=self.config.warmup_iters,
-            max_iters=self.config.max_iters,
+            cosine_schedule_period_iters=self.config.cosine_schedule_period_iters,
             lr=self.config.learning_rate,
             weight_decay=self.config.weight_decay,
             out_writer=self.writer,

tests/conftest.py

@@ -193,7 +193,7 @@ def tiny_config(tmp_path):
         "n_layers": 1,
         "train_label_smoothing": 0.01,
         "warmup_iters": 1,
-        "max_iters": 1,
+        "cosine_schedule_period_iters": 1,
         "max_epochs": 20,
         "val_check_interval": 1,
         "model_save_folder_path": str(tmp_path),

tests/unit_tests/test_config.py

@@ -37,3 +37,15 @@ def test_override(tmp_path, tiny_config):
 
     with pytest.raises(KeyError):
         Config(filename)
+
+
+def test_deprecated(tmp_path, tiny_config):
+    filename = str(tmp_path / "config_deprecated.yml")
+    with open(tiny_config, "r") as f_in, open(filename, "w") as f_out:
+        cfg = yaml.safe_load(f_in)
+        # Insert deprecated config option.
+        cfg["max_iters"] = 1
+        yaml.safe_dump(cfg, f_out)
+
+    with pytest.warns(DeprecationWarning):
+        Config(filename)