SquimObjectiveBundle

class torchaudio.pipelines.SquimObjectiveBundle

Data class that bundles associated information to use pretrained SquimObjective model.

This class provides interfaces for instantiating the pretrained model along with the information necessary to retrieve pretrained weights and additional data to be used with the model.

Torchaudio library instantiates objects of this class, each of which represents a different pretrained model. Client code should access pretrained models via these instances.

This bundle can estimate objective metric scores for speech enhancement, such as STOI, PESQ, Si-SDR. A typical use case would be a flow like waveform -> list of scores. Please see below for the code example.

Example: Estimate the objective metric scores for the input waveform.

>>> import torch
>>> import torchaudio
>>> from torchaudio.pipelines import SQUIM_OBJECTIVE as bundle
>>>
>>> # Load the SquimObjective bundle
>>> model = bundle.get_model()
Downloading: "https://download.pytorch.org/torchaudio/models/squim_objective_dns2020.pth"
100%|████████████| 28.2M/28.2M [00:03<00:00, 9.24MB/s]
>>>
>>> # Resample audio to the expected sampling rate
>>> waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate)
>>>
>>> # Estimate objective metric scores
>>> scores = model(waveform)
>>> print(f"STOI: {scores[0].item()},  PESQ: {scores[1].item()}, SI-SDR: {scores[2].item()}.")

Properties

sample_rate

property SquimObjectiveBundle.sample_rate

Sample rate of the audio that the model is trained on.

Type:

float

Methods

get_model

SquimObjectiveBundle.get_model(*, dl_kwargs=None) → SquimObjective

Construct the SquimObjective model, and load the pretrained weight.

The weight file is downloaded from the internet and cached with torch.hub.load_state_dict_from_url()

Parameters:

dl_kwargs (dictionary of keyword arguments) – Passed to torch.hub.load_state_dict_from_url().

Returns:

Variation of SquimObjective.