From a7497d749ef8828e0ce830dbb0d9923577ec92e6 Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Tue, 30 Jun 2026 11:21:24 +0200 Subject: [PATCH 1/7] make function a method and improve performance --- ms2deepscore/models/SiameseSpectralModel.py | 135 ++++++++++++++------ 1 file changed, 95 insertions(+), 40 deletions(-) diff --git a/ms2deepscore/models/SiameseSpectralModel.py b/ms2deepscore/models/SiameseSpectralModel.py index af42ab0c..e311505a 100644 --- a/ms2deepscore/models/SiameseSpectralModel.py +++ b/ms2deepscore/models/SiameseSpectralModel.py @@ -1,8 +1,9 @@ import os -from typing import Optional, Union, Dict, Any +from typing import Optional, Union, Dict, Any, Literal from pathlib import Path import numpy as np +import torch from torch import save, cat, zeros, cuda, no_grad from torch import device as torch_device from torch.nn.functional import relu @@ -75,6 +76,88 @@ def save(self, filepath: Union[str, Path]) -> None: # Important: no custom objects outside tensors/strings/primitives. save(checkpoint, str(filepath)) + def compute_embedding_array( + self, + spectra, + datatype: Literal["numpy", "pytorch"] = "numpy", + device: Optional[torch.device | str] = None, + batch_size: int = 1024, + progress_bar: bool = True, + ): + """ + Compute embeddings for a list of matchms Spectrum objects. + + Parameters + ---------- + spectra: + List of spectra to embed. + datatype: + "numpy" returns a NumPy array. + "pytorch" returns a CPU torch tensor. + device: + Device used for inference. If None, CUDA is used when available. + batch_size: + Number of spectra processed per encoder forward pass. + progress_bar: + Show progress bar. + + Returns + ------- + np.ndarray or torch.Tensor + Embedding array with shape (n_spectra, embedding_dim). + """ + datatype = datatype.lower() + if datatype not in {"numpy", "pytorch"}: + raise ValueError("datatype can only be 'numpy' or 'pytorch'.") + + if batch_size <= 0: + raise ValueError("batch_size must be a positive integer.") + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + device = torch.device(device) + + n_spectra = len(spectra) + embedding_dim = self.model_settings.embedding_dim + + self.to(device) + self.eval() + + if datatype == "numpy": + embeddings = np.empty((n_spectra, embedding_dim), dtype=np.float32) + else: + embeddings = torch.empty((n_spectra, embedding_dim), dtype=torch.float32) + + batch_starts = range(0, n_spectra, batch_size) + + with no_grad(): + for start in tqdm( + batch_starts, + total=(n_spectra + batch_size - 1) // batch_size, + desc="Computing spectral embeddings ...", + disable=not progress_bar, + ): + stop = min(start + batch_size, n_spectra) + batch_spectra = spectra[start:stop] + + spectra_tensors, metadata_tensors = tensorize_spectra( + batch_spectra, + self.model_settings, + ) + + batch_embeddings = self.encoder( + spectra_tensors.to(device), + metadata_tensors.to(device), + ).detach().cpu() + + if datatype == "numpy": + embeddings[start:stop, :] = batch_embeddings.numpy() + else: + embeddings[start:stop, :] = batch_embeddings + + return embeddings + class PeakBinner(nn.Module): """ @@ -351,45 +434,17 @@ def compute_embedding_array( datatype="numpy", device=None, progress_bar: bool = True, - ): + batch_size: int = 1024): """ - Compute the embeddings of all given spectra (list of matchms Spectrum objects). + Compatibility wrapper. - Parameters - ---------- - model: - A trained SiameseSpectralModel used to compute spectral embeddings. - spectra: - A list (or other iterable) of spectra to be embedded. - datatype: - Determines the output type of the embedding array: - - "numpy": returns a NumPy array of shape (n_spectra, embedding_dim). - - "pytorch": returns a PyTorch tensor of shape (n_spectra, embedding_dim). - device: - The device on which to perform the computation. - If None, it automatically uses CUDA if available, otherwise CPU. - progress_bar: - Whether to display a progress bar during embedding computation. + Prefer: + model.compute_embedding_array(...) """ - if datatype.lower() not in ["numpy", "pytorch"]: - raise ValueError("datatype can only be 'numpy' or 'pytorch'.") - if datatype.lower() == "numpy": - embeddings = np.zeros((len(spectra), model.model_settings.embedding_dim)) - else: - embeddings = zeros((len(spectra), model.model_settings.embedding_dim)) - - if device is None: - device = torch_device("cuda" if cuda.is_available() else "cpu") - model.to(device) - for i, spec in tqdm( - enumerate(spectra), - total=len(spectra), - desc="Computing spectral embeddings ...", - disable=not progress_bar): - X = tensorize_spectra([spec], model.model_settings) - with no_grad(): - if datatype.lower() == "numpy": - embeddings[i, :] = model.encoder(X[0].to(device), X[1].to(device)).cpu().detach().numpy() - else: - embeddings[i, :] = model.encoder(X[0].to(device), X[1].to(device)).cpu().detach() - return embeddings + return model.compute_embedding_array( + spectra, + datatype=datatype, + device=device, + progress_bar=progress_bar, + batch_size=batch_size, + ) From cd381c64fca54f4b83bce899722c5381f7bd0ff2 Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Tue, 30 Jun 2026 11:21:33 +0200 Subject: [PATCH 2/7] add batch_size parameter --- ms2deepscore/MS2DeepScore.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/ms2deepscore/MS2DeepScore.py b/ms2deepscore/MS2DeepScore.py index 477cae0a..e550140e 100644 --- a/ms2deepscore/MS2DeepScore.py +++ b/ms2deepscore/MS2DeepScore.py @@ -55,8 +55,14 @@ def __init__(self, model: SiameseSpectralModel, progress_bar: bool = True): self.output_vector_dim = self.model.model_settings.embedding_dim self.progress_bar = progress_bar - def get_embedding_array(self, spectrums): - return compute_embedding_array(self.model, spectrums, progress_bar=self.progress_bar) + def get_embedding_array(self, spectrums, batch_size: int = 1024) -> np.ndarray: + """Calculate the spectrum embeddings for a list of spectrums.""" + return compute_embedding_array( + self.model, + spectrums, + progress_bar=self.progress_bar, + batch_size=batch_size + ) def pair(self, reference: Spectrum, query: Spectrum) -> float: """Calculate the MS2DeepScore similaritiy between a reference and a query spectrum. From 9bc4929d6b4050b6625f3e4c2e2422d604e0363b Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Tue, 30 Jun 2026 11:28:26 +0200 Subject: [PATCH 3/7] allow passing batch_size and datatype --- ms2deepscore/MS2DeepScore.py | 3 ++- ms2deepscore/MS2DeepScoreEvaluated.py | 9 +++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/ms2deepscore/MS2DeepScore.py b/ms2deepscore/MS2DeepScore.py index e550140e..1bfa0773 100644 --- a/ms2deepscore/MS2DeepScore.py +++ b/ms2deepscore/MS2DeepScore.py @@ -55,11 +55,12 @@ def __init__(self, model: SiameseSpectralModel, progress_bar: bool = True): self.output_vector_dim = self.model.model_settings.embedding_dim self.progress_bar = progress_bar - def get_embedding_array(self, spectrums, batch_size: int = 1024) -> np.ndarray: + def get_embedding_array(self, spectrums, datatype: str = "numpy", batch_size: int = 1024) -> np.ndarray: """Calculate the spectrum embeddings for a list of spectrums.""" return compute_embedding_array( self.model, spectrums, + datatype=datatype, progress_bar=self.progress_bar, batch_size=batch_size ) diff --git a/ms2deepscore/MS2DeepScoreEvaluated.py b/ms2deepscore/MS2DeepScoreEvaluated.py index 3e5b33bb..3cba4137 100644 --- a/ms2deepscore/MS2DeepScoreEvaluated.py +++ b/ms2deepscore/MS2DeepScoreEvaluated.py @@ -69,8 +69,13 @@ def __init__(self, model: SiameseSpectralModel, self.output_vector_dim = self.model.model_settings.embedding_dim self.progress_bar = progress_bar - def get_embedding_array(self, spectrums, datatype="numpy"): - return compute_embedding_array(self.model, spectrums, datatype) + def get_embedding_array(self, spectrums, datatype="numpy", batch_size=1024): + return self.model.compute_embedding_array( + spectrums, + datatype=datatype, + batch_size=batch_size, + progress_bar=self.progress_bar, + ) def get_embedding_evaluations(self, embeddings): """Compute the RMSE. From 75ccdab3ca283fdc37fa53f16f1e293e68e873cc Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Tue, 30 Jun 2026 11:29:44 +0200 Subject: [PATCH 4/7] expand tests --- tests/test_ms2deepscore.py | 8 +++- tests/test_siamese_spectra_model.py | 64 +++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/tests/test_ms2deepscore.py b/tests/test_ms2deepscore.py index 9c9d081e..18b83d35 100644 --- a/tests/test_ms2deepscore.py +++ b/tests/test_ms2deepscore.py @@ -22,11 +22,15 @@ def get_test_ms2deepscore_instance(): return spectrums, model, similarity_measure -def test_MS2DeepScore_vector_creation(): +@pytest.mark.parametrize("batch_size", [16, None]) +def test_MS2DeepScore_vector_creation(batch_size): """Test embeddings creation. """ spectrums, _, similarity_measure = get_test_ms2deepscore_instance() - embeddings = similarity_measure.get_embedding_array(spectrums) + if batch_size is None: + embeddings = similarity_measure.get_embedding_array(spectrums) + else: + embeddings = similarity_measure.get_embedding_array(spectrums, batch_size=batch_size) assert embeddings.shape == (76, 100), "Expected different embeddings shape" assert isinstance(embeddings, np.ndarray), "Expected embeddings to be numpy array" diff --git a/tests/test_siamese_spectra_model.py b/tests/test_siamese_spectra_model.py index 88b225da..7fbbca64 100644 --- a/tests/test_siamese_spectra_model.py +++ b/tests/test_siamese_spectra_model.py @@ -1,5 +1,6 @@ import numpy as np import pytest +import torch from matchms import Spectrum from ms2deepscore.models.SiameseSpectralModel import (SiameseSpectralModel, train) @@ -151,3 +152,66 @@ def test_model_training(simple_training_spectra): # Check if bias in data is handled correctly assert (np.array(history["collection_targets"]) == 1).sum() == 200 assert (np.array(history["collection_targets"]) < .2).sum() == 200 + + +def test_siamese_model_compute_embedding_array_batched(dummy_spectra): + """Test that batched embedding computation matches single-spectrum batching.""" + settings = SettingsMS2Deepscore( + mz_bin_width=1.0, + base_dims=(100,), + embedding_dim=20, + train_binning_layer=False, + ) + model = SiameseSpectralModel(settings) + + embeddings_batch_size_1 = model.compute_embedding_array( + dummy_spectra, + datatype="numpy", + batch_size=1, + progress_bar=False, + ) + embeddings_batch_size_2 = model.compute_embedding_array( + dummy_spectra, + datatype="numpy", + batch_size=2, + progress_bar=False, + ) + + assert isinstance(embeddings_batch_size_1, np.ndarray) + assert isinstance(embeddings_batch_size_2, np.ndarray) + assert embeddings_batch_size_1.shape == (len(dummy_spectra), settings.embedding_dim) + assert embeddings_batch_size_2.shape == (len(dummy_spectra), settings.embedding_dim) + assert np.allclose(embeddings_batch_size_1, embeddings_batch_size_2, atol=1e-6) + + embeddings_torch = model.compute_embedding_array( + dummy_spectra, + datatype="pytorch", + batch_size=2, + progress_bar=False, + ) + + assert isinstance(embeddings_torch, torch.Tensor) + assert embeddings_torch.shape == (len(dummy_spectra), settings.embedding_dim) + assert embeddings_torch.device.type == "cpu" + assert not embeddings_torch.requires_grad + assert np.allclose( + embeddings_batch_size_2, + embeddings_torch.numpy(), + atol=1e-6, + ) + + +def test_siamese_model_compute_embedding_array_invalid_datatype(dummy_spectra): + settings = SettingsMS2Deepscore( + mz_bin_width=1.0, + base_dims=(100,), + embedding_dim=20, + ) + model = SiameseSpectralModel(settings) + + with pytest.raises(ValueError, match="datatype"): + model.compute_embedding_array( + dummy_spectra, + datatype="invalid", + progress_bar=False, + ) From 60a334d3b80de7915826701dc74132d9a2bee469 Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Tue, 30 Jun 2026 11:37:18 +0200 Subject: [PATCH 5/7] linting --- ms2deepscore/MS2DeepScoreEvaluated.py | 3 +-- ms2deepscore/models/SiameseSpectralModel.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/ms2deepscore/MS2DeepScoreEvaluated.py b/ms2deepscore/MS2DeepScoreEvaluated.py index 3cba4137..f4f6f615 100644 --- a/ms2deepscore/MS2DeepScoreEvaluated.py +++ b/ms2deepscore/MS2DeepScoreEvaluated.py @@ -4,8 +4,7 @@ from matchms.similarity.BaseSimilarity import BaseSimilarity from ms2deepscore.models.LinearEmbeddingEvaluation import \ compute_error_predictions -from ms2deepscore.models.SiameseSpectralModel import (SiameseSpectralModel, - compute_embedding_array) +from ms2deepscore.models.SiameseSpectralModel import (SiameseSpectralModel) from ms2deepscore.vector_operations import (cosine_similarity, cosine_similarity_matrix) diff --git a/ms2deepscore/models/SiameseSpectralModel.py b/ms2deepscore/models/SiameseSpectralModel.py index e311505a..3cda1823 100644 --- a/ms2deepscore/models/SiameseSpectralModel.py +++ b/ms2deepscore/models/SiameseSpectralModel.py @@ -4,8 +4,7 @@ import numpy as np import torch -from torch import save, cat, zeros, cuda, no_grad -from torch import device as torch_device +from torch import save, cat, no_grad from torch.nn.functional import relu from torch.optim import Adam from torch import nn From 319f58877931f605f48c5a3e40669da07d7a4431 Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Tue, 30 Jun 2026 11:50:32 +0200 Subject: [PATCH 6/7] linting, small refactor --- ms2deepscore/MS2DeepScore.py | 6 ++---- ms2deepscore/MS2DeepScoreEvaluated.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/ms2deepscore/MS2DeepScore.py b/ms2deepscore/MS2DeepScore.py index 1bfa0773..8f7bdf7a 100644 --- a/ms2deepscore/MS2DeepScore.py +++ b/ms2deepscore/MS2DeepScore.py @@ -2,8 +2,7 @@ import numpy as np from matchms import Spectrum from matchms.similarity.BaseSimilarity import BaseSimilarity -from ms2deepscore.models.SiameseSpectralModel import (SiameseSpectralModel, - compute_embedding_array) +from ms2deepscore.models.SiameseSpectralModel import SiameseSpectralModel from .vector_operations import cosine_similarity, cosine_similarity_matrix @@ -57,8 +56,7 @@ def __init__(self, model: SiameseSpectralModel, progress_bar: bool = True): def get_embedding_array(self, spectrums, datatype: str = "numpy", batch_size: int = 1024) -> np.ndarray: """Calculate the spectrum embeddings for a list of spectrums.""" - return compute_embedding_array( - self.model, + return self.model.compute_embedding_array( spectrums, datatype=datatype, progress_bar=self.progress_bar, diff --git a/ms2deepscore/MS2DeepScoreEvaluated.py b/ms2deepscore/MS2DeepScoreEvaluated.py index f4f6f615..de9f5f8b 100644 --- a/ms2deepscore/MS2DeepScoreEvaluated.py +++ b/ms2deepscore/MS2DeepScoreEvaluated.py @@ -4,7 +4,7 @@ from matchms.similarity.BaseSimilarity import BaseSimilarity from ms2deepscore.models.LinearEmbeddingEvaluation import \ compute_error_predictions -from ms2deepscore.models.SiameseSpectralModel import (SiameseSpectralModel) +from ms2deepscore.models.SiameseSpectralModel import SiameseSpectralModel from ms2deepscore.vector_operations import (cosine_similarity, cosine_similarity_matrix) From 6099e98ad18571270b78f8524d8d6cd3a5af25a1 Mon Sep 17 00:00:00 2001 From: Florian Huber Date: Tue, 30 Jun 2026 11:50:47 +0200 Subject: [PATCH 7/7] make progress bar count spectra (not batches) --- ms2deepscore/models/SiameseSpectralModel.py | 45 +++++++++++---------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/ms2deepscore/models/SiameseSpectralModel.py b/ms2deepscore/models/SiameseSpectralModel.py index 3cda1823..f02e6881 100644 --- a/ms2deepscore/models/SiameseSpectralModel.py +++ b/ms2deepscore/models/SiameseSpectralModel.py @@ -131,29 +131,32 @@ def compute_embedding_array( batch_starts = range(0, n_spectra, batch_size) with no_grad(): - for start in tqdm( - batch_starts, - total=(n_spectra + batch_size - 1) // batch_size, + with tqdm( + total=n_spectra, desc="Computing spectral embeddings ...", + unit="spectrum", disable=not progress_bar, - ): - stop = min(start + batch_size, n_spectra) - batch_spectra = spectra[start:stop] - - spectra_tensors, metadata_tensors = tensorize_spectra( - batch_spectra, - self.model_settings, - ) - - batch_embeddings = self.encoder( - spectra_tensors.to(device), - metadata_tensors.to(device), - ).detach().cpu() - - if datatype == "numpy": - embeddings[start:stop, :] = batch_embeddings.numpy() - else: - embeddings[start:stop, :] = batch_embeddings + ) as progress: + for start in batch_starts: + stop = min(start + batch_size, n_spectra) + batch_spectra = spectra[start:stop] + + spectra_tensors, metadata_tensors = tensorize_spectra( + batch_spectra, + self.model_settings, + ) + + batch_embeddings = self.encoder( + spectra_tensors.to(device), + metadata_tensors.to(device), + ).detach().cpu() + + if datatype == "numpy": + embeddings[start:stop, :] = batch_embeddings.numpy() + else: + embeddings[start:stop, :] = batch_embeddings + + progress.update(stop - start) return embeddings