A-SEM: Accelerated Speech Enhancement Metrics

A high-performance PyTorch library for computing speech quality metrics with GPU acceleration. Includes optimized implementations of PESQ, STOI, SDR, LSD, DNSMOS, and SpeechBERTScore.

Installation

Either download it directly through pip

pip install asem

or clone this repo and run (make sure uv is installed):

uv sync
source .venv/bin/activate

Usage

import torch
from asem import PESQ, STOI, SDR, LSD, DNSMOS, SpeechBERTScore

# Load your audio (shape: [batch_size, samples])
# 4 samples, 10 second each at 16kHz
clean_speech = torch.randn(4, 160000)
noisy_speech = torch.randn(4, 160000)

# Initialize metrics
pesq = PESQ(sample_rate=16000, use_gpu=True)
stoi = STOI(sample_rate=16000, use_gpu=True)
sdr = SDR(sample_rate=16000, use_gpu=True)

# Compute metrics
pesq_scores = pesq(clean_speech, noisy_speech)
stoi_scores = stoi(clean_speech, noisy_speech)
sdr_scores = sdr(clean_speech, noisy_speech)

print(pesq_scores)  # [{'PESQ': 2.1}, {'PESQ': 1.8}, ...]
print(stoi_scores)  # [{'STOI': 0.85, 'ESTOI': 0.82}, ...]

To work with variable-length utterances, pad the waveforms to a common length and pass the true sample counts via the optional lengths argument:

from torch.nn.utils.rnn import pad_sequence

clean_utterances = [torch.randn(120000), torch.randn(90000), torch.randn(143200)]
noisy_utterances = [torch.randn_like(x) for x in clean_utterances]

clean_batch = pad_sequence(clean_utterances, batch_first=True)
noisy_batch = pad_sequence(noisy_utterances, batch_first=True)
lengths = torch.tensor([x.shape[-1] for x in clean_utterances])

pesq_scores = pesq(clean_batch, noisy_batch, lengths=lengths)

Performance

Our GPU-accelerated implementations provide significant speedups over existing libraries:

while maintaining results that are extremely close to the originals

Available Metrics

Benchmarking

To run benchmarks on your system:

python benchmark_metrics.py
python plot_results.py

Acknowledgments

The PESQ implementation is based on the excellent work by audiolabs/torch-pesq.