A PyTorch implementation of a U-Net deep learning model for speech denoising using 2D spectrograms. This project is developed as part of the Machine Learning for Vision and Multimedia course at Politecnico di Torino.
ML-SpeechDenoising provides a complete pipeline for training and evaluating neural speech enhancement models. By employing spectral masking, the model learns to predict a magnitude mask that, when multiplied with a noisy input spectrogram, recovers the underlying clean speech.
- ML-SpeechDenoising
- U-Net Architecture: Optimized encoder-decoder model for spectral denoising.
- Multi-Dataset Support: Integration with LibriSpeech, Musan, DNS Challenge, Speech Noise Dataset by Haydar Kadioglu, and custom datasets.
- Unified Dataset Adapter: Integration with LibriSpeech, MUSAN, DNS Challenge, and HK-Dataset through a consistent adapter interface.
- Schema-Driven Configuration: Pydantic-validated
config.yamlensures mathematical consistency (e.g., FFT size vs. U-Net depth). - Advanced Training: Support for Multi-Scale Spectral Loss, Automatic Mixed Precision (AMP), Early Stopping, and top-K Checkpoint Management.
- CLI Interface: Full project management (downloading, training, evaluation) via a robust Typer-based CLI.
- Quality Metrics: Built-in support for objective metrics (PESQ, STOI, SDR).
- TensorBoard Integration: Real-time monitoring of loss, metrics, audio samples, and spectrograms.
This project uses uv for fast, reliable dependency management.
# Clone the repository
git clone <repository-url>
cd ML-SpeechDenoising
# Create a virtual environment and install all dependencies
uv syncGenerate a default config.yaml file to get started:
uv run ml-sd configDownload the required datasets (LibriSpeech, MUSAN, etc.):
uv run ml-sd data download allStart the training process:
uv run ml-sd trainEvaluate a trained model on a specific dataset:
uv run ml-sd eval models/model.ptDenoise an arbitrary audio file:
uv run ml-sd denoise input.wav --model models/model.pt --out output.wavML-SpeechDenoising/
├── src/ # Core implementation
│ ├── config/ # Pydantic schemas and configuration loading
│ ├── data/ # Dataset adapters and mixing
│ ├── model/ # U-Net architecture and custom layers
│ ├── training/ # Training engine, losses, and optimizers
│ ├── evaluation/ # Inference and audio metrics
│ └── utils/ # Signal processing, logging, and system utilities
├── scripts/ # CLI entry points
├── datasets/ # Raw and processed audio data (ignored by git)
├── models/ # Checkpoints and exported models (ignored by git)
├── logs/ # TensorBoard events and log files (ignored by git)
├── tests/ # Comprehensive pytest suite
└── notebooks/ # Data exploration and results analysis
The config.yaml file allows fine-grained control over the project. Below is an example of the core parameters with their descriptions:
system:
device: auto # Preferred compute device: 'cuda', 'mps', 'cpu', or 'auto'
num_workers: 4 # Number of worker processes for data loading
random_seed: 42 # Global random seed for reproducibility
signal:
sample_rate: 16000 # Project-wide sample rate (Hz)
n_fft: 512 # Window size for STFT
hop_length: 256 # Step size for STFT (overlap)
target_frames: 256 # Fixed time-length for training samples (approx 4s at 16kHz)
snr_range_db: # Random SNR range for mixing clean speech and noise
- -5.0
- 20.0
model:
base_features: 32 # Number of features in the first U-Net layer
depth: 4 # Number of U-Net levels (downsampling/upsampling steps)
training:
batch_size: 32 # Number of samples per training batch
learning_rate: 0.0001 # Initial learning rate for the optimizer
epochs: 100 # Maximum number of training epochs
optimizer: adamw # Optimizer type
patience: 10 # Early stopping patience (epochs to wait for improvement)
loss:
type: multi-scale # Multi-scale spectral loss for better reconstruction
evaluation:
eval_dataset_name: hkdataset # Dataset used for the final evaluation
batch_size: 1 # Evaluation batch size (usually 1 for varied lengths)The project uses an Adapter Pattern to unify different dataset formats into a common clean/noise pipeline.
| Dataset | Split | Description |
|---|---|---|
| LibriSpeech | Train / Val | High-quality clean speech. |
| MUSAN | Train / Val | Diverse noise (technical, ambient, music) for mixing. |
| DNS-Challenge | Evaluation (Unused) | Microsoft's Deep Noise Suppression challenge dataset. |
| HK-Dataset | Evaluation | Speech Noise Dataset by Haydar Kadioglu (Noisy/Clean pairs). |
Use ml-sd data download <name> to fetch specific datasets.
- Dataset Download: Fetch clean speech and noise from public sources (LibriSpeech, MUSAN).
- Data Caching: Filter and cache valid audio files for faster loading.
- Audio Mixing: Dynamically combine clean speech and noise at a random SNR (-5dB to 20dB).
- Augmentation: Apply spectral masking, frequency warping, and spectral scaling to improve generalization.
- Training Loop: Train the U-Net using Multi-Scale Spectral Loss (combining magnitude and log-magnitude L1/L2).
- Validation: Monitor performance on a validation set every epoch.
- Checkpointing: Automatically save the top-K best models based on validation loss.
- Final Evaluation: Test the model on the HK-Dataset for objective quality reporting.
The core of the project is a Spectral U-Net designed for 1-channel magnitude spectrograms.
- Encoder: A series of convolutional blocks that downsample the spectrogram while increasing feature depth, capturing hierarchical spectral patterns.
- Bottleneck: A latent space that represents the most salient features of the speech.
- Decoder: Mirroring the encoder, it upsamples the features back to the original spectrogram resolution.
- Skip Connections: Concatenate encoder features directly to the decoder to preserve high-frequency details.
- Output Layer: A Sigmoid activation predicts a Magnitude Mask (values between 0 and 1).
Input: [B, 1, Freq, Time] (noisy magnitude spectrogram)
↓
Encoder (4 levels: DoubleConv + MaxPool2d)
↓
Bottleneck (DoubleConv)
↓
Decoder (4 levels: Upsample + Concatenate + DoubleConv)
↓
Output: [B, 1, Freq, Time] (magnitude mask [0, 1])
- Base Features: 32 (doubles at each encoder level)
- Depth: 4 (downsampling/upsampling steps)
- Activation:
LeakyReLU(0.2) - Normalization:
BatchNorm2d - Output Layer:
Sigmoid(predicts the noise-filtering mask)
- DoubleConvBlock: (Conv2d 3x3 + BatchNorm + LeakyReLU) × 2. captures hierarchical spectral patterns.
- UpConvBlock: ConvTranspose2d (upsampling) followed by concatenation with skip connections and a DoubleConvBlock to preserve high-frequency details.
We use industry-standard objective metrics to evaluate speech quality and intelligibility:
| Metric | Range | Best Value | Meaning |
|---|---|---|---|
| PESQ (Perceptual Evaluation of Speech Quality) | -0.5 to 4.5 | 4.5 | Estimates human speech quality. |
| STOI (Short-Time Objective Intelligibility) | 0 to 1 | 1.0 | Measures how easy the denoised speech is for a human to understand. |
| SDR (Signal-to-Distortion Ratio) | -∞ to ∞ (dB) | High | Quantifies the overall audio quality relative to distortions and noise. |
These metrics are handled as follows:
- During Training: Calculated for a single validation sample per epoch and logged to TensorBoard alongside audio and spectrograms for qualitative monitoring.
- During Evaluation: Calculated for the entire test dataset, providing both baseline and enhanced scores. Results are saved as CSV and JSON reports in the
eval_outputs/directory.
Monitoring training is essential for detecting overfitting or convergence issues.
- Metrics: Loss (Train/Val), PESQ, STOI, and SDR are logged every epoch.
- Audio Samples: Real-time audio clips of (Noisy, Clean, Denoised) signals are uploaded periodically.
- Spectrograms: Visual comparison of the masks and reconstructed signals.
To launch the dashboard:
uv run tensorboard --logdir logs/tensorboardFor interactive analysis and visualization, the project includes:
notebooks/01_data_exploration.ipynb: Dataset analysis and audio statisticsnotebooks/02_results_analysis.ipynb: Results visualization and model comparison
To start the Jupyter interface:
uv run jupyter notebook notebooks/Run the full pytest suite to ensure system integrity:
# Run all tests
uv run pytest tests/
# Run specific test
uv run pytest tests/test_model.py -v