LiSenNet / README.md
claroche1's picture
conv-hardened: add frame-by-frame streaming graphs (fp32 + signed int8, 17 FIFO states); document both NPU deploy paths (windowed 1.15 ms/frame PESQ 2.998; streaming 16 ms hop PESQ 2.982)
4fe3356 verified
|
Raw
History Blame Contribute Delete
8.47 kB
metadata
license: mit
library_name: pytorch
tags:
  - speech-enhancement
  - audio
  - denoising
  - onnx
  - causal
  - streaming
  - real-time
  - edge-ai
  - stm32
datasets:
  - JacobLinCool/VoiceBank-DEMAND-16k

LiSenNet

Ultra-compact, causal, real-time speech enhancers trained on VoiceBank-DEMAND-16k — a sub-band U-Net with a magnitude-only mask (phase from a 2-iteration Griffin-Lim offline, or the noisy phase for real-time). Port of Yan, Zhou, Chen & Lu, LiSenNet, arXiv:2409.13285 (hyyan2k/LiSenNet, MIT).

This repo holds four variants, each in its own subfolder:

subfolder recipe params NPU-compiles FP32 PESQ real-time int8 PESQ
gru/ dual-path GRU (faithful) 36,783 3.006 2.930
conv/ dual-path conv 41,063 2.970 2.855
conv-hardened/ conv + NPU-hardened 36,288 3.013 2.998
conv-hardened-deep/ hardened + deep RF + ReLU6 46,248 ✓* 3.084 3.014 (int8) / 3.052 (hybrid)

PESQ is wideband, on the full 824-utterance VoiceBank-DEMAND test split. (*) conv-hardened-deep/ uses the same op set as conv-hardened/ plus Clip (ReLU6); its graph has not yet been through a stedgeai compile, conv-hardened/ has (topology verified on Neural-ART).

  • gru/ is the faithful reproduction and the original quality reference. Its GRU + 2-axis LayerNorm do not compile to the STM32N6 Neural-ART NPU.
  • conv/ replaces the GRU bottleneck with a dual-path conv one (0 GRU / 0 LayerNormalization). Its ops map to the NPU, but the FIFO-state streaming graph (conv/g_best_streaming_fp32.onnx, feat + N state_i_in -> est_mag + N state_i_out) crashes the Neural-ART codegen — kept as the CPU/onnxruntime frame-by-frame reference.
  • conv-hardened/ is the compile-verified NPU-deployable variant: per-channel BatchNorm (folds into the convs), ReLU, plain ConvTranspose upsampling, and a stateless windowed deploy graph (conv-hardened/g_best_windowed_int8_static.onnx, signed QInt8, feat_window (B,3,132,257) -> est_mag (B,64,257), window = receptive field 68 + 64 emitted frames) that compiles to Neural-ART — the artifact handed to stedgeai. The hardened primitives also quantize far better (int8 drop −0.016 vs −0.115 for conv/). It also ships a frame-by-frame streaming graph — see below.
  • conv-hardened-deep/ is the best model overall: the hardened recipe, deeper (3 blocks) with an extra dilation stage (receptive field 196 frames ≈ 3.1 s) and ReLU6 activations (bounded ranges quantize better; exports as Clip). Window is 196+64=260 frames (feat_window (B,3,260,257)). It ships two signed windowed int8 artifacts: g_best_windowed_int8_static.onnx (everything int8, PESQ 3.014) and g_best_windowed_int8_decoder_fp32.onnx (int8 except the decoder's QDQ nodes, PESQ 3.052 — the int8 loss is decoder-localized; the decoder then runs as float epochs on the board).

Two deploy paths (conv-hardened/) — throughput vs latency

The hardened variant ships both NPU-deployable graphs. They are the same trained weights; they differ only in how much audio each inference call consumes. Both are measured on an STM32N6570-DK (STM32N657, MCU 800 MHz / NPU 1 GHz, stedgeai 4.0.1); PESQ is the full 824-utterance VBD test split, int8 + noisy phase.

graph shape algorithmic latency latency/frame RTF PESQ
g_best_windowed_int8_static.onnx feat_window (B,3,132,257) -> est_mag (B,64,257) 1.02 s block 1.15 ms 0.072 2.998
g_best_streaming_int8_static.onnx feat (B,3,1,257) + 17 states -> est_mag (B,1,257) + 17 states one 16 ms hop 2.79 ms 0.174 2.982
  • the windowed graph is stateless and the throughput champion (1.15 ms per emitted frame), but buffers 1.02 s of audio per call — use it when block latency is acceptable;
  • the streaming graph is true frame-by-frame — one 16 ms hop in, one enhanced frame out, with 17 bounded FIFO state tensors carried on-device (47 KiB weights
    • 147 KiB activations, entirely on-chip). It costs ~2.4× more compute per frame but drops algorithmic latency by ~64×, which makes it the best-quality streaming speech enhancer measured on this NPU.

Both are signed QInt8 (the Neural-ART recipe). The streaming graph is calibrated by threading the real FIFO state through the trained model, so the state tensors see realistic ranges; its int8 cost is −0.018 PESQ vs the same-protocol FP32 streaming reference (3.000).

Code + full write-up: https://github.com/LarocheC/eco8-neaixt — see RESULTS_LISENNET.md.

Files (per subfolder)

config.json, g_best (PyTorch {"generator": state_dict}), g_best_fp32.onnx and g_best_int8_static.onnx (whole-utterance mask sub-network, feat (B,3,T,F) -> est_mag (B,T,F)). conv/ additionally has g_best_streaming_fp32.onnx and g_best_streaming_int8_static.onnx (single frame + explicit state I/O — CPU/onnxruntime reference only; this variant's streaming graph does not compile to Neural-ART). conv-hardened/ has both NPU deploy graphs: g_best_windowed_{fp32,int8_static}.onnx (stateless windowed) and g_best_streaming_{fp32,int8_static}.onnx (frame-by-frame, 17 FIFO states) — see the two-deploy-paths table above. The ONNX graphs are the mask sub-network only — STFT, feature build and phase recovery stay host-side.

Loading (PyTorch)

import json, torch
from huggingface_hub import hf_hub_download
from common.env import AttrDict
from lisennet.model import build_lisennet

REPO, SUB = "claroche1/LiSenNet", "conv-hardened"      # or "gru" / "conv"
cfg  = json.load(open(hf_hub_download(REPO, f"{SUB}/config.json")))
ckpt = torch.load(hf_hub_download(REPO, f"{SUB}/g_best"), map_location="cpu", weights_only=True)
model = build_lisennet(AttrDict(cfg)).eval()
model.load_state_dict(ckpt["generator"])   # model(noisy_wav)["est"]

Running the NPU windowed deploy graph (conv-hardened/)

Stateless: feed a sliding window of the last 68 + 64 = 132 feature frames and read the 64 newest enhanced-magnitude frames (no state tensors to carry).

import numpy as np, onnxruntime as ort
from huggingface_hub import hf_hub_download

sess = ort.InferenceSession(
    hf_hub_download("claroche1/LiSenNet", "conv-hardened/g_best_windowed_int8_static.onnx"),
    providers=["CPUExecutionProvider"],
)
feat_window = np.zeros((1, 3, 132, 257), np.float32)   # last 68+64 feature frames
est_mag = sess.run(["est_mag"], {"feat_window": feat_window})[0]  # (1, 64, 257)

Running the streaming graph frame-by-frame (conv-hardened/, conv/)

One 16 ms hop in, one enhanced frame out; the FIFO state tensors are threaded from each call's outputs into the next call's inputs (start-of-stream = zeros). Point it at conv-hardened/ for the NPU-deployable graph (PESQ 2.982), or conv/ for the CPU-only reference.

import numpy as np, onnxruntime as ort
from huggingface_hub import hf_hub_download

sess = ort.InferenceSession(
    hf_hub_download("claroche1/LiSenNet", "conv-hardened/g_best_streaming_int8_static.onnx"),
    providers=["CPUExecutionProvider"],
)
state_in = [i for i in sess.get_inputs() if i.name != "feat"]   # FIFO states
out_names = [o.name for o in sess.get_outputs()]                # est_mag + state_*_out
zeros = lambda s: np.zeros([d if isinstance(d, int) else 1 for d in s], np.float32)
states = {i.name: zeros(i.shape) for i in state_in}            # start-of-stream = zeros

def step(feat_t):                                              # feat_t: (1, 3, 1, 257)
    res = sess.run(out_names, {"feat": feat_t, **states})
    for i, v in zip(state_in, res[1:]):
        states[i.name] = v
    return res[0]                                              # est_mag (1, 1, 257)

License

MIT. See the source repository for training code and full attribution.