diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml new file mode 100644 index 0000000..0eb8f5d --- /dev/null +++ b/.github/workflows/python-tests.yml @@ -0,0 +1,24 @@ +name: Python Tests + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install . + - name: Run tests + run: pytest diff --git a/README.md b/README.md index 2b44b23..0d3177d 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ The MAEST models are also available for inference from [Essentia](https://essent # Install Our software has been tested in Ubuntu 22.04 LTS and CentOS 7.5 using Python 3.10. -If Python 3.10 is unavailable in your environment, we recommend using the [Conda](https://docs.conda.io) package manager to set up the working environment. +If Python 3.10 is unavailable in your environment, we recommend using the [Conda](https://docs.conda.io) package manager to set up the working environment. 1. Create a conda environment (optional): @@ -34,7 +34,7 @@ pip install -e . # Usage -We use [Sacred](https://github.com/IDSIA/sacred) to run, configure and log our experiments. +We use [Sacred](https://github.com/IDSIA/sacred) to run, configure and log our experiments. The different routines can be run with Sacred commands, and many experiment options can be directly configure from the command line. @@ -115,7 +115,7 @@ the paper: - `extract_embeddings`: returns a [3, 768] vector with the embeddings for each audio file in the `predict` dataset. The three dimensions of the first axis correspond to the CLS token, the - DIST token and the average of the rest of tokens (see Section 4.1 from the paper). + DIST token and the average of the rest of tokens (see Section 4.1 from the paper). - `extract_logits`: returns logits that can be used in the teacher student setup, or transformed into label predictions by applying a `Sigmoid` function. Each pre-training configuration has its inference version. For example, to extract embeddings @@ -173,7 +173,8 @@ _, embeddings = model(audio, transformer_block=6) MAEST is designed to accept `data` in different input formats: - 1D: 16kHz audio waveform is assumed. -- 2D: mel-spectrogram is assumed (frequency, time). +- 2D: (with `melspectrogram_input=False`) audio is assumed (batch, time). +- 2D: (`melspectrogram_input=True`) mel-spectrogram is assumed (frequency, time). - 3D: batched mel-spectrogram (batch, frequency, time). - 4D: batched mel-spectrgroam plus singleton channel axis (batch, 1, frequency, time). @@ -190,7 +191,6 @@ The following `arch` values are supported: - `discogs-maest-30s-pw-73e-ts` - `discogs-maest-30s-pw-129e-519l` - Additionally, `predict_labels()` is an auxiliary function that applies a sigmoid activation, averages the predictions along the time axes, and returns the label vector for convenience. ```python @@ -213,4 +213,3 @@ If you are going to use MAEST as part of your research, please consider citing t year={2023}, } ``` - diff --git a/models/helpers/melspectrogram.py b/models/helpers/melspectrogram.py new file mode 100644 index 0000000..c6e3719 --- /dev/null +++ b/models/helpers/melspectrogram.py @@ -0,0 +1,60 @@ +import torch +from torch.nn import Module +from torchaudio.transforms import Spectrogram, MelScale + +# According to Pytoch, mel-spectrogram should be implemented as a module: +# https://pytorch.org/audio/stable/transforms.html + +# WARNING: The torchaudio implementation is similar but not identical to Essentia's. +# Relative tolerance is 1e-3 and absolute tolerance is 1e-3. We assume that this has minimal +# impact in the resulting embeddings. + + +class MelSpectrogram(Module): + """Extract mel-spectgrams as a torchaudio module""" + + sr = 16000 + win_len = 512 + hop_len = 256 + power = 2 + n_mel = 96 + norm = "slaney" + mel_scale_type = "slaney" + norm_mean = 2.06755686098554 + norm_std = 1.268292820667291 + + def __init__(self): + super().__init__() + + self.spec = Spectrogram( + n_fft=self.win_len, + win_length=self.win_len, + hop_length=self.hop_len, + power=self.power, + ) + + self.mel_scale = MelScale( + n_mels=self.n_mel, + sample_rate=self.sr, + n_stft=self.win_len // 2 + 1, + norm=self.norm, + mel_scale=self.mel_scale_type, + ) + + def znorm(self, input_values: torch.Tensor) -> torch.Tensor: + return (input_values - (self.norm_mean)) / (self.norm_std * 2) + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + # convert to power spectrogram + spec = self.spec(waveform) + + # convert to mel-scale + mel = self.mel_scale(spec) + + # apply logC compression + logmel = torch.log10(1 + mel * 10000) + + # normalize + logmel = self.znorm(logmel) + + return logmel diff --git a/models/helpers/melspectrogram_extractor.py b/models/helpers/melspectrogram_extractor.py deleted file mode 100644 index 0f9d380..0000000 --- a/models/helpers/melspectrogram_extractor.py +++ /dev/null @@ -1,56 +0,0 @@ -import numpy as np -from essentia.streaming import ( - VectorInput, - MonoLoader, - FrameCutter, - TensorflowInputMusiCNN, -) -from essentia import Pool, run, reset - - -class MelSpectrogramExtractor: - # mel-spec signature params - sample_rate = 16000 - resample_quality = 4 - frame_size = 512 - hop_size = 256 - - def __init__(self): - self.pool = Pool() - self.frame_cutter = FrameCutter( - frameSize=self.frame_size, - hopSize=self.hop_size, - silentFrames="keep", - ) - self.melspec = TensorflowInputMusiCNN() - - self.frame_cutter.frame >> self.melspec.frame - self.melspec.bands >> (self.pool, "mel_bands") - - def __call__(self, audio): - # make sure audio is float32 for Essentia - audio = audio.astype(np.float32) - - # set vector input and connect the network - vector_input = VectorInput(audio) - vector_input.data >> self.frame_cutter.signal - - # run the network and get the mel bands - run(vector_input) - mel_bands = np.array(self.pool["mel_bands"]) - - # clear the pool and reset the vector input - reset(vector_input) - vector_input.data.disconnect(self.frame_cutter.signal) - self.pool.clear() - del vector_input - - # return freq, time - return mel_bands.T - - def load_audio(self, audio_file): - return MonoLoader( - filename=audio_file, - sampleRate=self.sample_rate, - resampleQuality=self.resample_quality, - )() diff --git a/models/helpers/vit_helpers.py b/models/helpers/vit_helpers.py index 2f6a45b..6b24e7b 100644 --- a/models/helpers/vit_helpers.py +++ b/models/helpers/vit_helpers.py @@ -2,12 +2,13 @@ Adapted from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py """ + import math import warnings from copy import deepcopy import torch -from timm.models.helpers import load_pretrained +from timm.models import load_pretrained from torch import nn diff --git a/models/maest.py b/models/maest.py index c2b876d..221026e 100644 --- a/models/maest.py +++ b/models/maest.py @@ -12,6 +12,7 @@ import warnings from collections import OrderedDict from functools import partial +from typing import Tuple import numpy as np import torch @@ -24,7 +25,7 @@ trunc_normal_, build_model_with_cfg, ) -from .helpers.melspectrogram_extractor import MelSpectrogramExtractor +from .helpers.melspectrogram import MelSpectrogram from .discogs_labels import discogs_400labels, discogs_519labels _logger = logging.getLogger("MAEST") @@ -242,7 +243,7 @@ def __init__( def forward(self, x): B, C, H, W = x.shape if not (H == self.img_size[0] and W == self.img_size[1]): - warnings.warn( + _logger.debug( f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." ) # to do maybe replace weights @@ -495,7 +496,6 @@ def __init__( ) self.num_tokens = 2 if distilled else 1 self.distilled_type = distilled_type - self.melspectrogram_extractor = None norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) act_layer = act_layer or nn.GELU if self.num_classes == 400: @@ -583,8 +583,7 @@ def __init__( self.init_weights(weight_init) - def init_melspectrogram_extractor(self): - self.melspectrogram_extractor = MelSpectrogramExtractor() + self.melspectrogram = MelSpectrogram() def init_weights(self, mode=""): assert mode in ("jax", "jax_nlhb", "nlhb", "") @@ -663,8 +662,9 @@ def forward_features(self, x, transformer_block=-1, return_self_attention=False) f"CUT time_new_pos_embed.shape: {time_new_pos_embed.shape}" ) else: - warnings.warn( - f"the patches shape:{x.shape} are larger than the expected time encodings {time_new_pos_embed.shape}, x will be cut" + raise Exception( + f"the patches shape:{x.shape} are larger than the expected time encodings {time_new_pos_embed.shape}," + " please reduce the input duration." ) x = x[:, :, :, : time_new_pos_embed.shape[-1]] x = x + time_new_pos_embed @@ -828,36 +828,55 @@ def forward_features(self, x, transformer_block=-1, return_self_attention=False) first_RUN = False return torch.cat([cls, dist, feats], dim=1) - def forward(self, x, transformer_block=-1, return_self_attention=False): + def forward( + self, + x, + transformer_block: int = -1, + return_self_attention: bool = False, + melspectrogram_input: bool = False, + ) -> Tuple[torch.Tensor | None, torch.Tensor]: + """ + Forward pass of the model. + Args: + x (torch.Tensor): Input tensor, either raw audio or melspectrogram. + transformer_block (int): If -1, returns the output of the last block. + If >= 0, returns the output of the specified block. + return_self_attention (bool): If True, returns self-attention from the specified block. + melspectrogram_input (bool): If True, input is expected to be a melspectrogram. + Returns: + activations (torch.Tensor): Output tensor with the model's predictions. Requires `transformer_block` to be -1. + embeddings (torch.Tensor): Output tensor with the model's embeddings of the selected transformer block. + """ + global first_RUN if first_RUN: _logger.debug(f"x size: {len(x)}") + assert isinstance(x, torch.Tensor), "Input must be a torch.Tensor" + assert x.nelement() > 0, "Input tensor must not be empty" + if len(x.shape) == 1: + assert melspectrogram_input is False, ( + "Input is 1D, but melspectrogram_input is True. This is not supported." + ) + _logger.debug("extracting melspec") - if not self.melspectrogram_extractor: - _logger.debug("initializing melspectrogram extractor") - self.init_melspectrogram_extractor() # extract melspec from raw audio - x = self.melspectrogram_extractor(x) - # normalize - x = (x - DISCOGS_MEAN) / (DISCOGS_STD * 2) - - # zero-pad to if if the input is shorter than the expected size - pad = self.img_size[1] - x.shape[1] - - if pad > 0.2 * self.img_size[1]: - warnings.warn( - "Padding input by more than 20% of the expected size. " - "This may result in poor performance. " - "Consider using a MAEST models trained with a smaller input size." - ) - - if pad > 0: - x = np.pad(x, ((0, 0), (0, pad)), mode="constant") + x = self.melspectrogram(x) + + if x.shape[1] >= self.img_size[1]: + # cut into equally-sized patches. + trim = x.shape[1] % self.img_size[1] + if trim: + x = x[:, :-trim] + x = x.reshape(self.img_size[0], 1, -1, self.img_size[1]) + # resort axes: batch, channels, freq, time + x = torch.swapaxes(x, 0, 2) + else: + x = x.reshape(1, 1, x.shape[0], x.shape[1]) - if len(x.shape) == 2: + elif len(x.shape) == 2 and melspectrogram_input: # need batching trim = x.shape[1] % self.img_size[1] if trim: @@ -867,7 +886,10 @@ def forward(self, x, transformer_block=-1, return_self_attention=False): # Note that the new batch axis needs to be next to the time. # resort axes: batch, channels, freq, time x = np.swapaxes(x, 0, 2) - x = torch.tensor(x) + + elif len(x.shape) == 2 and not melspectrogram_input: + x = self.melspectrogram(x) + x.unsqueeze_(1) x = self.forward_features( x, diff --git a/pyproject.toml b/pyproject.toml index cba7484..5058421 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,17 +4,16 @@ build-backend = "setuptools.build_meta" [project] name = "maest" -version = "0.1.0" +version = "0.1.1" description = "MAEST model package" authors = [ { name = "Pablo Alonso", email = "pablo.alonso@upf.edu" } ] license = { text = "AGPL-3.0 license" } readme = "README.md" -requires-python = ">=3.7" +requires-python = ">=3.10" dependencies = [ - "essentia", "timm~=0.9", "torchmetrics~=0.11.4", "tqdm~=4.65", @@ -24,6 +23,8 @@ dependencies = [ "tensorboard~=2.13.0", "numpy~=1.26", "torch", + "torchaudio", + "pytest", ] [project.urls] @@ -31,7 +32,7 @@ homepage = "https://github.com/palonso/maest" repository = "https://github.com/palonso/maest" [tool.setuptools] -packages = ["maest", "maest.helpers"] +packages = ["maest", "maest.helpers"] [tool.setuptools.package-dir] "maest" = "models" diff --git a/tests/test_maest.py b/tests/test_maest.py new file mode 100644 index 0000000..8d81d75 --- /dev/null +++ b/tests/test_maest.py @@ -0,0 +1,57 @@ +import pytest +import numpy as np +import torch +from maest import get_maest + + +@pytest.fixture +def model(): + # Use a simple config for testing; adjust as needed for your environment + return get_maest(arch="discogs-maest-30s-pw-129e", pretrained=False) + + +def test_numpy_input(model): + input_data = np.random.rand(128, 128) + with pytest.raises(Exception): + model(input_data) + + +def test_empty_input(model): + input_data = torch.empty([]) + with pytest.raises(Exception): + model(input_data) + + +def test_1d_input(model): + input_data = torch.rand(10 * 16000).float() + activations, _ = model(input_data) + assert activations.shape == (1, 400), ( + "Activations shape should be (batch_size, num_classes)" + ) + + +def test_2d_audio_activations(model): + input_data = torch.rand(2, 10 * 16000).float() + # Batch of 2 audio samples + activations, _ = model(input_data, melspectrogram_input=False) + assert activations.shape == (2, 400), ( + "Output batch size should match input batch size" + ) + + +def test_2d_melspec_activations(model): + input_data = torch.rand(96, 1875).float() + # Batch of 10 mel spectrograms + activations, _ = model(input_data, melspectrogram_input=True) + assert activations.shape == (1, 400), ( + "Output batch size should match input batch size" + ) + + +def test_2d_melspec_embeddings(model): + input_data = torch.rand(96, 1875).float() + # Batch of 10 mel spectrograms + _, embeddings = model(input_data, melspectrogram_input=True, transformer_block=6) + assert embeddings.shape == (1, 2304), ( + "Output batch size should match input batch size" + )