Modernize optimizer

ms2deepscore/SettingsMS2Deepscore.py

@@ -86,7 +86,7 @@ class SettingsMS2Deepscore:
         base_dims:
             The in between layers to be used. Default = (2000, 2000, 2000)
         embedding_dim:
-            The dimension of the final embedding. Default = 400
+            The dimension of the final embedding. Default = 300
         ionisation_mode:
             The ionisation mode that is used for training the model.
         balanced_sampling_across_ionmodes:
@@ -97,6 +97,9 @@ class SettingsMS2Deepscore:
             The dropout rate that should be used during training
         learning_rate:
             The learning rate that should be used during training.
+        weight_decay:
+            The weight decay that should be used during training (optimizer is AdamW).
+            Default is 0.01, which is a common value for AdamW. If you want to use Adam, set weight_decay=0.0.
         epochs:
             The number of epochs that should be used during training.
         patience:
@@ -185,8 +188,8 @@ def __init__(self, validate_settings=True, **settings):
         self.train_test_split_fraction = 20
 
         # model structure
-        self.base_dims = (10000,)
-        self.embedding_dim = 500
+        self.base_dims = (2000, 2000, 2000)
+        self.embedding_dim = 300
         self.ionisation_mode = "positive"
         self.activation_function = "relu"
         self.balanced_sampling_across_ionmodes = False
@@ -199,8 +202,9 @@ def __init__(self, validate_settings=True, **settings):
         # training settings
         self.dropout_rate = 0.0
         self.learning_rate = 0.00025
+        self.weight_decay = 0.01
         self.epochs = 250
-        self.patience = 20
+        self.patience = 10
         self.loss_function = "mse"
         self.weighting_factor = 0
 

ms2deepscore/models/SiameseSpectralModel.py

@@ -10,7 +10,7 @@
 from torch.export.dynamic_shapes import Dim
 from torch.nn.functional import relu
 from torch.onnx import export
-from torch.optim import Adam
+from torch.optim import AdamW
 from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
 from ms2deepscore.__version__ import __version__
@@ -354,6 +354,7 @@ def train(
     data_generator,
     num_epochs: int,
     learning_rate: float,
+    weight_decay: float = 0.0,
     validation_loss_calculator=None,
     early_stopping=True,
     patience: int = 10,
@@ -379,8 +380,10 @@ def train(
         Number of epochs for training.
     learning_rate
         Learning rate for the optimizer.
-    val_generator (iterator, optional)
-        An iterator for validation data batches.
+    weight_decay
+        Weight decay for the optimizer (AdamW optimizer). Default is 0.0.
+    validation_loss_calculator
+        An object that calculates validation loss. If None, no validation is performed.
     early_stopping
         Whether to use early stopping.
     patience
@@ -407,7 +410,13 @@ def train(
         raise ValueError(f"Unknown loss function. Must be one of: {LOSS_FUNCTIONS.keys()}")
     criterion = LOSS_FUNCTIONS[loss_function.lower()]
 
-    optimizer = Adam(model.parameters(), lr=learning_rate)
+    optimizer = AdamW(
+        model.parameters(),
+        lr=learning_rate,
+        weight_decay=weight_decay,
+        foreach=False,
+        fused=False,
+        )
 
     # Initialize TensorBoard writer
     if checkpoint_filename:

tests/test_SettingsMS2deepscore.py

@@ -80,7 +80,7 @@ def test_save_settings(tmp_path):
         result = json.load(file)
     assert result["epochs"] == 200
     assert result["base_dims"] == [200, 200]
-    assert result["embedding_dim"] == 500
+    assert result["embedding_dim"] == 300
 
 
 def test_get_dict():
@@ -89,7 +89,7 @@ def test_get_dict():
     settings_dict = settings.get_dict()
     assert settings_dict["epochs"] == 200
     assert settings_dict["base_dims"] == (200, 200)
-    assert settings_dict["embedding_dim"] == 500
+    assert settings_dict["embedding_dim"] == 300
 
 
 def test_get_dict_with_coerced_types():