moonbox-dev/MossFormerGAN_SE_16K-ONNX

MossFormerGAN ONNX

ONNX builds of MossFormerGAN, made by Moonbox for on-device use with ONNX Runtime (CPU, CUDA, DirectML). Each precision is one folder; what the graphs take and return, and the signal processing that goes around them, is under How to run.

Files

The last column is the benchmark on the VoiceBank-DEMAND test set (824 noisy utterances; details under How to run): each precision's speech-quality scores, and how far its output is from the original PyTorch model's.

Precision	File	Size	sha256	Benchmark
fp16	fp16/model.onnx	12.2 MB	4047ebd56656	PESQ 3.342, STOI 0.9570, SI-SDR 19.77 dB; vs fp32 59.8 / 57.2 dB; worst file -0.05 dB (CUDA)
int8	int8/model.onnx	11.2 MB	3fbe341fb6ed	PESQ 3.344, STOI 0.9570, SI-SDR 19.75 dB; vs fp32 44.4 / 41.3 dB; worst file -0.16 dB (CUDA; CPU within 0.01 dB of it on 20 files)

How to run

MossFormerGAN SE 16K (ClearerVoice-Studio's speech enhancement: removes noise from speech at 16 kHz; a MossFormer generator trained with a metric discriminator) as one ONNX graph: the generator that turns the compressed noisy spectrum into the enhanced one. Input spec [1, 2, T, 201] (real and imaginary parts of a 400-point STFT, hop 100, Hamming window, power-law compressed by 0.3), outputs real and imag [1, 1, T, 201]. T is free, up to 8192 frames (51 s) per call. The normalization, the STFT, the compression and the inverse STFT stay outside the graph, exactly as in ClearerVoice's own code. So the simplest correct pipeline is ClearerVoice's decode function with its network swapped for this graph:

import numpy as np, onnxruntime as ort, soundfile as sf, torch
import clearvoice.utils.decode as dec                      # pip install clearvoice==0.1.2 (only for its decode helper)
from clearvoice.network_wrapper import network_wrapper

w = network_wrapper(); w.model_name = "MossFormerGAN_SE_16K"; w.load_args_se()
args = w.args; args.network = "MossFormerGAN_SE_16K"
sess = ort.InferenceSession("fp16/model.onnx", providers=["CUDAExecutionProvider", "CPUExecutionProvider"])

class Net:                                                  # what decode() calls: net(spec) -> (real, imag)
    def __call__(self, spec):
        return [torch.from_numpy(o) for o in sess.run(None, {"spec": spec.numpy().astype(np.float32)})]
    def eval(self):
        return self

noisy, sr = sf.read("noisy_16k.wav", dtype="float32")      # 16 kHz mono
clean = dec.decode_one_audio(Net(), torch.device("cpu"), noisy[None], args)
sf.write("enhanced.wav", np.asarray(clean, np.float32).reshape(-1)[: len(noisy)], 16000)

decode_one_audio itself processes files longer than 10 s in 10 s windows. Attention runs over the whole window, so memory grows with the square of its length: about 1 GB per attention step on a 10 s clip. On CUDA, let ONNX Runtime shrink its arena after each call (run option memory.enable_memory_arena_shrinkage = gpu:0), or its buffers pile up across clips of different lengths.

Precisions. fp16/: weights and math in fp16, inputs and outputs fp32, for GPU providers. int8/: weight-only int8 (every MatMul, Gemm and Conv weight int8 per output channel, activations fp32). It runs on the CPU provider, but there the model is about 7 times slower than real time on a laptop CPU (the PyTorch original: 4 times), so use a GPU for anything but short clips. No int4: 4-bit MatMul weights put 459 of the 824 test files more than 0.5 dB below the original (worst 3.5 dB), and the file would be only 0.6 MB smaller than int8/. No bf16 (no bf16 convolution in ONNX Runtime). The 12 attention modules' identical rotary tables are stored once.

Benchmark. VoiceBank-DEMAND test set (Valentini et al.: 824 utterances by two unseen speakers), resampled to 16 kHz, enhanced with ClearerVoice's own decode code. PESQ is wide-band (ITU-T P.862.2), STOI and SI-SDR against the clean speech. The PyTorch original (clearvoice 0.1.2, fp32) runs 4 times slower than real time on this laptop's CPU, so it was run on the first 50 files. There, the ONNX fp32 graph on CUDA matches it at 77.3 dB mean (75.6 dB worst file), and SI-SDR differs by at most 0.002 dB. The full set then compares against that fp32 graph. "vs fp32" is the SI-SDR of each output against the fp32 graph's, mean / worst file; "worst file" is the largest SI-SDR drop on any one file. Pass (the speech-enhancement threshold used across these repositories): no file more than 0.5 dB below. Real-time factor on an RTX 3080 Laptop GPU, or with 4 CPU threads (Intel laptop CPU) for the CPU rows, whole pipeline.

	Files	PESQ	STOI	SI-SDR (dB)	vs fp32 (dB)	Worst file	Real-time factor
noisy input	824	1.97	0.921	8.39	-	-	-
original (PyTorch fp32, CPU)	first 50	3.442	0.9485	20.81	-	-	3.9
ONNX fp32 (CUDA, not published)	first 50	3.442	0.9485	20.81	-	-0.00 dB vs original	-
ONNX fp32 (CUDA, not published)	824	3.342	0.9570	19.77	-	-	0.18
fp16/ (CUDA)	824	3.342	0.9570	19.77	59.8 / 57.2	-0.05 dB	0.20
int8/ (CUDA)	824	3.344	0.9570	19.75	44.4 / 41.3	-0.16 dB	0.18
int8/ (CPU)	first 20	3.514	0.9710	21.19	43.6 / 41.1	-0.08 dB	7.0
int4 (not published)	824	3.334	0.9560	19.02	28.0 / 20.1	-3.53 dB (459 files below -0.5)	0.18

How these were made

From MossFormerGAN at revision dcb1b20e1e, starting from an fp32 ONNX graph (where it came from is said above). fp16: onnxconverter-common with inputs and outputs kept fp32, or, where the network's activations leave fp16's range, fp16 weights with fp32 math (said above when so). int8: weight-only, every Conv / MatMul / GRU weight int8 with a per-output-channel scale behind DequantizeLinear and activations in fp32, so it runs on the CPU provider without a calibration set. int4: ONNX Runtime's MatMulNBits (weight-only, block 32) on the MatMul weights, the rest int8 as above, where the graph has enough MatMul weights to matter. No bf16: ONNX Runtime has no bf16 convolution or GRU kernels. Every graph was benchmarked before upload; the sha256 column above is what the files must hash to.

License

Apache-2.0. These files are a derivative of the original model and carry its license: it is copied in this repository, and it governs your use of these files. The original model card follows unchanged.

---

Original model card

The MossFormerGAN_SE_16K model weights for 16 kHz speech enhancement in ClearerVoice-Studio repo.

This model is trained on large scale datasets inclduing open-sourced and private data.

It enhances speech audios by removing background noise.