Missing src/data/datamodules and src/data/transforms

[DBraun]

Would you mind adding src/data/datamodules and transforms to the repo to make reproducibility a bit easier? I would like to train on MusicCaps just to debug some things. Hopefully the code you have is not too locked down due toe PrivateCaps. Thanks!

[nateraw]

Hey there @DBraun - thanks for your work on dawdreamer :)

I've been hacking in new datasets by adding new file to src/data/your_data.py containing different dataset + datamodule and pointing to the "target" in a new hydra config file @ configs/data/your_config.yaml. Either way, you'll want to precompute features...can then wrap with your own dataset or make csv compatible with audio_csv.py pointing to the saved features.

For you + those lurking...I think a minimal example to compute feats is like this (only diff is I use torchaudio here + use correct hparams for the checkpoint):

# ! pip install torchlibrosa
import random

import torch
import torchaudio
import torchaudio.functional as F
from torch import nn
from torch.nn.functional import pad
from torchlibrosa.stft import LogmelFilterBank, Spectrogram


class SpectrogramExtractor(nn.Module):
    def __init__(
        self,
        mel_bins=64,
        sample_rate=48000,
        spec_window_size=1024,
        hop_size=480,
        fmin=50,
        fmax=14000,
    ):
        super().__init__()
        window = "hann"
        center = True
        pad_mode = "reflect"
        ref = 1.0
        amin = 1e-10
        top_db = None
        self.spectrogram_extractor = Spectrogram(
            n_fft=spec_window_size,
            hop_length=hop_size,
            win_length=spec_window_size,
            window=window,
            center=center,
            pad_mode=pad_mode,
            freeze_parameters=True,
        )
        self.spectrogram_extractor.double()
        # Logmel feature extractor
        self.logmel_extractor = LogmelFilterBank(
            sr=sample_rate,
            n_fft=spec_window_size,
            n_mels=mel_bins,
            fmin=fmin,
            fmax=fmax,
            ref=ref,
            amin=amin,
            top_db=top_db,
            freeze_parameters=True,
        )
        self.logmel_extractor.double()

    def forward(self, x):
        dtype = x.dtype
        x = x.double()
        x = self.spectrogram_extractor(x)  # (batch_size, 1, time_steps, freq_bins)

        x = self.logmel_extractor(x)  # (batch_size, 1, time_steps, mel_bins)

        x = x.type(dtype)
        x = x.transpose(2, 3).squeeze(0)
        return x


def load_audio(audio_fpath: str, target_sr: int, target_frames: int = None, mono=True):
    y, sr = torchaudio.load(audio_fpath)
    if mono:
        y = y.mean(dim=0)
    n_samples = y.shape[0]

    if sr != target_sr:
        y = F.resample(y, sr, target_sr)
        sr = target_sr
        n_samples = y.shape[0]

    if target_frames:
        if n_samples < target_frames:
            y = pad(y, (0, target_frames - n_samples))
        elif n_samples > target_frames:
            start_idx = random.randint(0, n_samples - target_frames - 1)
            y = y[start_idx : start_idx + target_frames]

    return y


if __name__ == '__main__':
    sr = 48000

    spec_extractor = SpectrogramExtractor(sample_rate=sr)
    audio_fpath = 'input.wav'

    wav = load_audio(audio_fpath, sr)
    spec = spec_extractor(wav.unsqueeze(0)).squeeze()
    print(f"{wav.shape=}")
    print(f"{spec.shape=}")

[aRI0U]

Hi!
As @nateraw mentioned you have the choice between using your own custom dataset class and using our audio_csv code. Tbh I think using our AudioCSVDatamodule/Dataset is the easiest solution.

To use the audio_csv classes, all you need is create a csv file (or several) with file paths and corresponding caption. By default, we assume that spectrograms are precomputed and stored as .npy files, which can significantly accelerate training depending on your infra.

You can otherwise load raw audio files using option audio_=True in AudioCSVDataset.
In that case the audio/spectrogram conversion will be done in the audio_encoder itself. This enables to do the conversion in parallel on the GPU directly. Therefore you'd have to make sure that precomputed_lms=False in the HTSAT class (src.networks.htsat.HTSATSwinTransformer).

Please let me know if there is something unclear of if your application seems incompatible with our proposed AudioCSV classes

[nateraw]

@aRI0U probably worth mentioning the AudioCSV dataset currently uses 16k sample rate and users would have to set that to 48k to use it correctly with pretrained ckpt