Speech Enhancement with Neural Network Quantization

Taking a compact speech-enhancement model from research to a quantization-ready state — with real measured numbers at every stage, from training through INT8 deployment.

PyTorch · ONNX Runtime · INT8 (PTQ + QAT) · GTCRN · VoiceBank+DEMAND

An end-to-end PyTorch → ONNX → INT8 pipeline for ultra-low-resource speech denoising, built around a faithful reimplementation of the GTCRN architecture (48,245 parameters) and extended with an original quantization pipeline (PTQ + QAT) that the reference codebase does not provide.

Scope note: The goal of this project is not state-of-the-art enhancement. It is a defensible, reproducible, end-to-end pipeline with real measured numbers — from training through quantization — including an honest treatment of where the open-source deployment stack falls short for recurrent architectures.

---

Results

Float32 baseline (VoiceBank+DEMAND test set, 824 utterances)

Metric	Achieved	Reference target
PESQ (wideband)	2.82	~2.87
STOI	0.941	~0.940

Deployable models — size, latency, quality

Latency is the mean of 50 inference runs after 10 warm-up runs, measured on CPU.

Model	Size (KB)	Latency (ms)	PESQ	STOI
PyTorch float32	589.1	191.7 ± 55.4	2.820	0.941
ONNX float32	392.2	72.8 ± 12.9	2.820	0.941
ONNX PTQ INT8	288.3	69.9 ± 2.9	2.679	0.940

Simulated full-model quantization (quality only — see Limitations)

Quantizes all layers including the GRUs, which off-the-shelf tooling cannot. Reported as simulated INT8 because the ONNX standard cannot serialize a quantized GRU.

Model	PESQ	STOI
PTQ-sim (all layers, no fine-tuning)	2.247	0.928
QAT-sim (all layers, fine-tuned)	2.666	0.936

Key result: under identical full-model quantization, QAT fine-tuning recovers ~73% of the PESQ gap that naive PTQ destroys (2.25 → 2.67, against a 2.82 baseline) — a direct validation of Jacob et al. (2018) on a speech-enhancement model.

Training curve

---

Architecture

The model is a faithful reimplementation of GTCRN (Rong et al., ICASSP 2024) — a time-frequency masking network engineered for ultra-low computational cost. Input is a 3-channel spectrogram (magnitude, real, imaginary); output is a complex ratio mask applied to the noisy spectrum.

Noisy waveform
    │  STFT (512 / 256 / 512, sqrt-Hann)
    ▼
3-channel feature (mag, real, imag)
    │  ERB compression (257 → 129 freq bins)
    │  Subband Feature Extraction
    ▼
Encoder  (Conv blocks + Grouped Temporal Conv blocks w/ dilated depthwise conv + TRA)
    │
    ▼
Dual-Path Grouped RNN ×2  (intra-frame + inter-frame grouped GRUs)
    │
    ▼
Decoder  (transpose Conv blocks, ERB decompression 129 → 257)
    │
    ▼
Complex ratio mask ── applied to noisy spectrum ──► Enhanced spectrum ──► iSTFT

Parameter count (48,245) is used throughout as a correctness checksum against the reference.

---

Pipeline

Stage	Script	Output
Dataset / STFT features	dataset.py	3-channel feature tensors
Model	model.py	SpeechEnhancer (48,245 params)
Loss (SI-SNR + mag + real/imag)	loss.py	hybrid loss
Training (float32)	train.py, kaggle_train.py	best.tar
Evaluation (PESQ / STOI)	evaluate.py	metrics, enhanced wavs
ONNX export	export.py	model.onnx
PTQ INT8 (deployable, ONNX Runtime)	quantize_ptq.py	model_ptq_int8.onnx
PTQ INT8 (simulated, all layers)	quantize_ptq_sim.py	metrics
QAT INT8 (simulated, all layers)	quantize_qat_sim.py, kaggle_qat.py	qat_best.tar, metrics
Fake-quant module (Jacob et al. scheme)	fake_quant.py	reusable quant wrappers
Benchmark (size / latency / quality)	benchmark.py	results/benchmark.txt

---

Setup

conda create -n speech_enhancement python=3.10
conda activate speech_enhancement
pip install -r requirements.txt

Download VoiceBank+DEMAND (a.k.a. VCTK-DEMAND) and arrange it as:

VoiceBank-DEMAND/
├── noisy_trainset_28spk_wav/
├── clean_trainset_28spk_wav/
├── noisy_testset_wav/
├── clean_testset_wav/
├── train.scp
└── test.scp

Usage

# Train (GPU recommended — use kaggle_train.py on Kaggle's free T4)
python train.py

# Evaluate the float32 model
python evaluate.py --backend torch

# Export to ONNX and verify fidelity
python export.py

# Evaluate the ONNX model (should match PyTorch exactly)
python evaluate.py --backend onnx

# Post-training quantization (deployable, ONNX Runtime)
python quantize_ptq.py
python evaluate.py --backend onnx --onnx-model checkpoints/model_ptq_int8.onnx

# Simulated full-model quantization studies
python quantize_ptq_sim.py
python quantize_qat_sim.py     # fine-tunes; run on GPU

# Full benchmark
python benchmark.py

---

Methodology notes

• Feature representation. STFT with a 512-point FFT, 256-hop, and square-root Hann window. The complex spectrogram is split into magnitude, real, and imaginary channels — magnitude carries spectral energy, real/imag preserve phase for the complex mask.
• Loss. Hybrid of time-domain SI-SNR (weight 0.01), compressed-magnitude MSE, and real/imag MSE, with magnitude/phase compression exponents of 0.3 / 0.7.
• Training. Adam (lr 1e-3), batch size 4, gradient clipping at 5.0, ReduceLROnPlateau (halve on 5-epoch plateau), 100 epochs, seeded for reproducibility. Converged by ~epoch 64.
• ONNX export. Uses the legacy TorchScript exporter with opset 18 to correctly honour the dynamic time axis; verified lossless (PyTorch vs ONNX PESQ/STOI identical to 4 d.p.).
• Quantization. Follows the integer-arithmetic scheme of Jacob et al. (2018). The fake-quant module (fake_quant.py) implements affine INT8 quantization with a straight-through estimator, wrapping Conv, Linear, and GRU layers. PTQ and QAT are compared under identical quantization coverage so the only differing variable is fine-tuning.

---

Limitations

• Recurrent layers are not natively quantizable by standard tooling. The architecture contains 14 GRU operators. ONNX Runtime has no INT8 GRU kernel and the ONNX standard defines no quantized GRU operator (QLinearGRU), so the deployable PTQ model quantizes only MatMul weights, leaving convolutions and GRUs in float32. This caps compression at 1.36×.
• QAT results are simulated INT8. The fully fake-quantized model measures the accuracy impact of full-model quantization but remains float32 on disk and is not, as-is, a deployable artifact.
• INT8 latency gains require hardware with integer arithmetic units. On CPU, INT8 improved latency stability (std 12.9 → 2.9 ms) more than raw speed.
• Utterance-level, not streaming. Causal padding was introduced during debugging as a step toward streaming, but a true streaming inference path is out of scope.

Future work

• GRU unrolling into elementary linear operations so the recurrent path exports as quantizable MatMul sequences (enabling genuine end-to-end INT8 deployment).
• Vendor toolchains: Qualcomm's AIMET for RNN-aware quantization and the QNN SDK for on-device Snapdragon benchmarking.
• Streaming inference and an accuracy-vs-size trade-off curve across architectural variants.

---

References

1. Rong et al., GTCRN: A Speech Enhancement Model Requiring Ultralow Computational Resources, ICASSP 2024. (official repo)
2. B. Jacob et al., Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference, CVPR 2018.
3. Valentini-Botinhao et al., VoiceBank+DEMAND (VCTK-DEMAND) noisy speech database.

License

MIT — see LICENSE.