diff --git a/ms2deepscore/MS2DeepScore.py b/ms2deepscore/MS2DeepScore.py index 477cae0a..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 @@ -55,8 +54,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, datatype: str = "numpy", batch_size: int = 1024) -> np.ndarray: + """Calculate the spectrum embeddings for a list of spectrums.""" + return self.model.compute_embedding_array( + spectrums, + datatype=datatype, + 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. diff --git a/ms2deepscore/MS2DeepScoreEvaluated.py b/ms2deepscore/MS2DeepScoreEvaluated.py index 3e5b33bb..de9f5f8b 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) @@ -69,8 +68,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. diff --git a/ms2deepscore/models/SiameseSpectralModel.py b/ms2deepscore/models/SiameseSpectralModel.py index af42ab0c..f02e6881 100644 --- a/ms2deepscore/models/SiameseSpectralModel.py +++ b/ms2deepscore/models/SiameseSpectralModel.py @@ -1,10 +1,10 @@ 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 -from torch import save, cat, zeros, cuda, no_grad -from torch import device as torch_device +import torch +from torch import save, cat, no_grad from torch.nn.functional import relu from torch.optim import Adam from torch import nn @@ -75,6 +75,91 @@ 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(): + with tqdm( + total=n_spectra, + desc="Computing spectral embeddings ...", + unit="spectrum", + disable=not progress_bar, + ) 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 + class PeakBinner(nn.Module): """ @@ -351,45 +436,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, + ) 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, + )