diff --git a/ms2deepscore/SettingsMS2Deepscore.py b/ms2deepscore/SettingsMS2Deepscore.py index 979b0a77..d1432b7b 100644 --- a/ms2deepscore/SettingsMS2Deepscore.py +++ b/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 diff --git a/ms2deepscore/models/SiameseSpectralModel.py b/ms2deepscore/models/SiameseSpectralModel.py index 25d768fb..89f9415d 100644 --- a/ms2deepscore/models/SiameseSpectralModel.py +++ b/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: diff --git a/tests/test_SettingsMS2deepscore.py b/tests/test_SettingsMS2deepscore.py index 532c71af..e43f1ef0 100644 --- a/tests/test_SettingsMS2deepscore.py +++ b/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():