Considerations for a speech enhancement model

[errorfixrepeat]

Firstly, thank you for providing such a useful library. I think it can only do good for the source separation and enhancement fields to have datasets, models, and recipes using shared terminology and building blocks.

I am particularly interested in high sample rate (44.1/48kHz) applications, and have been trying to adapt the existing models with varying success. For my application, both the speech and the noise signals are of interest. As a starting point I've decided to recreate a model by Wichern and Lukin which seems relatively simple Paper.

My understanding of it:
Encoder: STFT (2048 Hann window, 512 stride)
Masker: The magnitude of the spectrum feeds 2 layers of 256 unit BLSTM to produce a speech magnitude mask with a sigmoid activation. The noise mask is the difference (1 - speech mask).
Decoder: The masked magnitudes are combined with the input phase, and then ISTFT back to audio.
Loss: MSE of the masked magnitude spectrum, plus a tunable crosstalk parameter.

One particular challenge is that I am unable to get separated training data in the target domain, so the speech and noise signals have been created using another separation model. I am doing some things to mitigate this:

• The training generator randomly samples noise files, so any speech residue in the noise signal (and vice versa) should be de-correlated with the target speech.
• I have added clean speech and noise sources, contributing about 2/3rd of the 60hr total training data.

What I've learnt:

• To use Asteroid loss functions outside of permutational invariant training you need to modify them to produce a single value (adding reduction='mean' to the loss function is the preferred method).
• Training in the encoded domain is significantly faster.
• I wasn't able to get Griffin-lim to work within the model. I think this was due to cuda/cpu location issues.
• Altering the model to use and predict complex masks provided an improvement in sound quality.
• Different maskers (e.g. ConvTas) with the same encoders didn't produce noticeably better results.
• I cannot get the models to overfit, even with all regularisation turned off.

What I'm looking to improve:
Restricting the noise mask to be the difference signal has improved the separation, but it's still not perfect. There seems to be a roughness in the speech output, and the noise output is still modulated by the speech to some extent. I have a few ideas of where to go, but I'd really welcome some thoughts and suggestions from the experts!

[mpariente]

Thanks for your detailed post.
Some follow-up questions:

To use Asteroid loss functions outside of permutational invariant training you need to modify them to produce a single value

Even the single_src functions?
Is the problem that you still have the batch dimension, or there is another one?
How could the user experience be improved for speech enhancement in your opinion?

Training in the encoded domain is significantly faster.

What do you mean exactly? Compared to having a time-domain loss? What the order of magnitude of "faster"?

I wasn't able to get Griffin-lim to work within the model. I think this was due to cuda/cpu location issues.

Could you find the error message, so that we can try to fix it. We don't have extensive testing on it, and it would be great if it were usable.

Altering the model to use and predict complex masks provided an improvement in sound quality.

Great! Which mask are you using exactly, and what is your loss function then? Do you still input only the magnitude to the masker?

Different maskers (e.g. ConvTas) with the same encoders didn't produce noticeably better results.

This is an unexpected result. I expect this might be due to feeding the magnitude only being fed to the network.
Also maybe to the SNR conditions, how close are you to oracle performance with your current model?

I cannot get the models to overfit, even with all regularisation turned off.

Is this true for all masker? Have you tried overfitting on a very small part of the data?

Regarding your dataset, I'm interested in the clean speech and noise generation. Have you validated that it helps the enhancement performance compared to using only the 40hrs of clean speech/noise pairs? If yes, do you have numbers?

FYI, I think there is some high sample rate noise here but not sure about its quality.

[errorfixrepeat]

Training in the encoded domain is significantly faster.

What do you mean exactly? Compared to having a time-domain loss? What the order of magnitude of "faster"?

As I am using a STFT MSE loss which is the same as my model, I can either:

• Run the full model, and then compute the spectrogram again in the loss function.
• Alter the model to return the spectrogram and use this directly in the loss function.

The former takes 21m50s for 1 epoch (including 1m03s validation), the latter is 6m38s (22s validation).
For reference, I am using a single 2080ti. Train is 35688 5 second 48kHz segments, validation is 3514, batch size is 12.

[mpariente]

So only the iSTFT and STFT before the loss and during the loss takes this long time?
What are the constant of the transform please?

[mpariente]

The intuitive solution would be for all multisrc losses to output a single value by default so they work independently, but I'm not sure how that would affect existing code.

That's impossible, because we need to return the loss for each pair, and each batch.

A 'non permutation invariant' wrapper would be simple, and easy to convey in the documentation.

That would be feasible, would you like to contribute this?

[jonashaag]

The former takes 21m50s for 1 epoch (including 1m03s validation), the latter is 6m38s (22s validation).

Can you show the exact code? Sounds like something is wrong there.

[errorfixrepeat]

So only the iSTFT and STFT before the loss and during the loss takes this long time?
What are the constant of the transform please?

Yes, I believe so. Parameters:
encoderparams={'fb_name':'stft', "n_filters":1024,"stride":256, "kernel_size":1024, "sample_rate":48000}

A 'non permutation invariant' wrapper would be simple, and easy to convey in the documentation.

That would be feasible, would you like to contribute this?

Sure, I'll give it a go.

[jonashaag]

FWIW I've been using Asteroid to train 24 kHz models for a while now and it works great. I've mostly been using DCUNet, DCCRN, CRUSE (not yet contributed to Asteroid).

Can you tell why you lack clean speech data? There are multiple open and free 44/48 kHz clean speech and also noise datasets.

[mpariente]

Is there a list of all high-sample-rate audio data somewhere?
It would be nice to have it actually :-)

[jonashaag]

We can start one, here are some starting points

• https://github.com/microsoft/DNS-Challenge/tree/master/datasets_fullband
• Some of those https://openslr.org/resources.php
• VCTK https://datashare.ed.ac.uk/handle/10283/2950
• DAPS https://ccrma.stanford.edu/~gautham/Site/daps.html
• Common Voice (has lots of noisy data as well) https://commonvoice.mozilla.org/en/datasets
• SWC is 48 kHz as well if I remember correctly https://nats.gitlab.io/swc/

[errorfixrepeat]

Lots to think about from yours and Manuel's responses, it may take me a little while to try your suggestions.

As you point out there is clean speech available, but my use case is differs in speaking style and microphone choice. The model I used to split the in-domain data was trained on some of the datasets above, and didn't perform optimally. This lead me to believe that I needed real data for training, but I haven't tested this assumption.

[jonashaag]

By the way, I believe all of your issues have nothing to do with sample rate, maybe we should adapt the title accordingly so as to not confuse future readers.

[errorfixrepeat]

I agree that other issues are present, so I've changed the title.

That said, I had no luck blindly applying the default models at 48kHz which were not due to data. These were primarily due to memory limitations, where the workarounds involve lowering the example duration and/or reducing the batch size. In addition the default parameters lead to a shorter contextual windows, which will ultimately hinder performance.

I think there's some room for discussion on which are more suitable, and to someone not well versed in the field it's not intuitive when certain architectures should be chosen over another.

Your earlier comment about experience training 24kHz models was useful, have you found any parameter changes were necessary to achieve good results?

[jonashaag]

The thing is that 48 kHz models aren’t researched as much as 16/8 kHz models so there aren’t a lot of results that we can implement here.

For 24 kHz the default params seem to work well. But for the stuff I’m working on (dereverberation) I think about the right data a lot more than the right model.

[errorfixrepeat]

I did a little more testing comparing the training time using TF output from the model and re-encoding in the loss function.

I based the code on the 00_Getting_Started notebook, adding the loss function and subclassing DPRNN to return TF representations. In addition I modified the Librimix getitem to repeat the duration of each example along the time axis by a global variable "reps". When reps is 1 each input is 40000, equating to 0.83* seconds at 48kHz.

# Asteroid is based on PyTorch and PyTorch-Lightning.
from torch import optim
from pytorch_lightning import Trainer

# We train the same model architecture that we used for inference above.
from asteroid.models import DPRNNTasNet

# In this example we use Permutation Invariant Training (PIT) and the SI-SDR loss.
from asteroid.losses import pairwise_neg_sisdr, PITLossWrapper

# MiniLibriMix is a tiny version of LibriMix (https://github.com/JorisCos/LibriMix),
# which is a free speech separation dataset.
from asteroid.data import LibriMix

# Asteroid's System is a convenience wrapper for PyTorch-Lightning.
from asteroid.engine import System

# Extra imports
from torch.nn.modules.loss import _Loss
from asteroid_filterbanks import transforms, make_enc_dec
from asteroid.utils.torch_utils import pad_x_to_y, script_if_tracing, jitable_shape
from asteroid.models.base_models import _shape_reconstructed, _unsqueeze_to_3d
import timeit
import pandas as pd
import matplotlib.pyplot as plt


# Define TF MSE loss, and overwrite DPRNN to return spectrograms
class FBMSE(_Loss):
    r"""Measure mean square error on a batch.
    Supports both tensors with and without source axis.

    Shape:
        - est_targets: :math:`(batch, ...)`.
        - targets: :math:`(batch, ...)`.

    Returns:
        :class:`torch.Tensor`: with shape :math:`(batch)`

    Examples
        >>> import torch
        >>> from asteroid.losses import PITLossWrapper
        >>> targets = torch.randn(10, 2, 32000)
        >>> est_targets = torch.randn(10, 2, 32000)
        >>> # singlesrc_mse / multisrc_mse support both 'pw_pt' and 'perm_avg'.
        >>> loss_func = PITLossWrapper(singlesrc_mse, pit_from='pw_pt')
        >>> loss = loss_func(est_targets, targets)
    """
    def __init__(self, reduction=None, encoder=None, transforms="labels", magnitude=True, scale=1):
        super().__init__()
        self.reduction = reduction
        if encoder:
            self.encoder = make_enc_dec(**encoder)[0]
        else:
            self.encoder = None

        self.transforms = transforms
        self.magnitude = magnitude
        self.scale =scale
        
    def forward(self, est_targets, targets):
        if self.transforms == "labels" and self.encoder:
            targets = self.encoder.forward(targets)
        elif self.transforms == "both" and self.encoder:
            targets = self.encoder.forward(targets)
            est_targets = self.encoder.forward(est_targets)
        if targets.size() != est_targets.size() or targets.ndim < 2:
            raise TypeError(
                f"Inputs must be of shape [batch, *], got {targets.size()} and {est_targets.size()} instead"
            )
        
        if self.magnitude:
            targets = transforms.magphase(targets)[0]
            est_targets = transforms.magphase(est_targets)[0]
        
        loss = (targets - est_targets) ** 2
        loss *= self.scale
        
        mean_over = list(range(1, loss.ndim))
        loss = loss.mean() if self.reduction == "mean" else loss.mean(dim=mean_over)
        return loss

class OverDPRNN(DPRNNTasNet):
    def __init__(self, n_src, return_spec = False, **encoderparams):
        super().__init__(n_src, **encoderparams)
        self.return_spec = return_spec

    def forward(self, wav):
        """Enc/Mask/Dec model forward
        Args:
            wav (torch.Tensor): waveform tensor. 1D, 2D or 3D tensor, time last.
        Returns
            torch.Tensor, of shape (batch, n_src, time) or (n_src, time).
        """
        # Remember shape to shape reconstruction, cast to Tensor for torchscript
        shape = jitable_shape(wav)
        # Reshape to (batch, n_mix, time)
        wav = _unsqueeze_to_3d(wav)

        # Real forward
        tf_rep = self.forward_encoder(wav)
        est_masks = self.forward_masker(tf_rep)
        masked_tf_rep = self.apply_masks(tf_rep, est_masks)
        if self.return_spec:
            return masked_tf_rep
        else:
            decoded = self.forward_decoder(masked_tf_rep)
            reconstructed = pad_x_to_y(decoded, wav)
            return _shape_reconstructed(reconstructed, shape)


# Hack to make fake input tensor at 48k
def new_getitem(self, index):
    item = self.__oldgetitem__(index)
    return item[0].repeat_interleave(reps, dim=-1), item[1].repeat_interleave(reps, dim=-1)

LibriMix.__oldgetitem__ = LibriMix.__getitem__
LibriMix.__getitem__ = new_getitem

# This will automatically download MiniLibriMix from Zenodo on the first run.
train_loader, val_loader = LibriMix.loaders_from_mini(task="sep_clean", batch_size=2, segment=5)

results = []
for reps in [1,2,4,8]:
    # Train DPRNN with time domain inputs to loss function.
    encoderparams={'fb_name':'stft', "n_filters":1024,"stride":256, "kernel_size":1024, "sample_rate":48000}
    model = OverDPRNN(n_src=2, **encoderparams)
    loss = FBMSE(reduction="mean", encoder = encoderparams, transforms="both")
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    system = System(model, optimizer, loss, train_loader, val_loader)
    trainer = Trainer(max_epochs=5, gpus=1)
    ttim = %timeit -n1 -r1 -o trainer.fit(system)

    # Clean up
    del model, loss, optimizer, system, trainer

    # Train model using TF inputs to loss function
    encoderparams={'fb_name':'stft', "n_filters":1024,"stride":256, "kernel_size":1024, "sample_rate":8000}
    model = OverDPRNN(return_spec = True, n_src=2, **encoderparams)
    loss = FBMSE(reduction="mean", encoder = encoderparams, transforms="labels")
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    system = System(model, optimizer, loss, train_loader, val_loader)
    trainer = Trainer(max_epochs=5, gpus=1)
    ftim = %timeit -n1 -r1 -o trainer.fit(system)

    # Clean up
    del model, loss, optimizer, system, trainer
    results.append((reps, ttim.all_runs[0], ftim.all_runs[0]))

df = pd.DataFrame(results, columns=["repeats","timedomain", "freqdomain"])
display(df)
plt.plot(df["repeats"], df["timedomain"], df["repeats"], df["freqdomain"],)

Below is the output, showing the time to train 5 epochs:

As you can see the training time is not a constant ratio between the direct TF and re-encoded model, becoming much worse with longer input lengths. I suspect it may be linked to unwrapping of recurrent layers, so plan on trying again with ConvTasNet once Colab allows me GPU access again (my local box is in use).

[errorfixrepeat]

Same behaviour occurs with ConvTasNet.

(now with correct screenshot)