Inference improvements

.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

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},
     }
 ```
-

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

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
 
 

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,

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 = "[email protected]" }
 ]
 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,14 +23,16 @@ dependencies = [
     "tensorboard~=2.13.0",
     "numpy~=1.26",
     "torch",
+    "torchaudio",
+    "pytest",
 ]
 
 [project.urls]
 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"