GTCRN-Light v3 Causal Stream

A causal version for real-time streaming inference. Converted from the offline version to causal, trading a small amount of performance for real-time capability.

This project is based on the work of zerong7777-boop.

Project link: zerong7777-boop/gtcrn-light

For the detailed design rationale, refer to [GTCRN 48 kHz] Evolution of the Causal-Stream Model

Note: v0 → v1 refers to the original GTCRN → GTCRN-Light process; my contribution there was minor. The main work here is the spectrum extension: 16kHz offline → 48kHz offline → 48kHz causal, which is also why the parameter count grew noticeably.

The project is currently being adapted for embedded deployment. The core operators + runtime memory have been compressed to under 500KB, but that portion of the code is not publicly available. I still hope this work can be useful to you.

I almost got fired because of this project. My boss doesn’t like me working on too much machine learning-related content, and I’m not sure yet whether I should continue developing it. But I’ve decided to make part of my current work public anyway, hoping it can help you.

Basic Info

• Best Epoch: 35
• DNSMOS_OVR: 2.983
• Parameters: 145,127
• Sample Rate: 48kHz
• Frame Latency: 10ms (single frame)

Differences from the Offline Version

The following changes were made to ensure causality:

1. GTConvLite → CausalGTConvLite: Left padding changed to (kernel-1) * dilation
2. TRALite → CausalTRA: Also changed to causal padding
3. DPGRNN inter: GRU on the time axis changed to unidirectional (frequency axis remains bidirectional, no causality impact)
4. Activation function: PReLU → SiLU
5. DSConv/DSDeconv: Added intermediate BN, order also adjusted

Performance Comparison

Version	Parameters	DNSMOS	Real-time
v1 Offline	139K	3.15	×
v2 Offline	139K	3.15	×
v3 Causal	145K	2.98	√

A drop of 0.17 points in exchange for real-time inference capability.

Network Architecture

Input spec (B, 513, T, 2)
    │
    ▼
ERB_48k.bm(): 513 → 219 bands
    │
    ▼
in_conv: Conv2d(2 → 3, k=1×1)
    │
    ▼
┌─ CausalEncoder ───────────────────────────────┐
│                                               │
│  DSConv (3→32ch, stride=2): 219 → 110        │
│      DWConv → BN → SiLU → PWConv → BN → SiLU │
│                                    ← skip1   │
│  DSConv (32→32ch, stride=2): 110 → 55        │
│                                    ← skip2   │
│                                               │
│  CausalGTConvLite × 6 (dilation: 1,2,4,8,4,2)│
│      Each layer: CausalDWConv(5×5) → BN → SiLU│
│                  → PWConv → BN → SiLU         │
│                  → CausalTRA → SE → Residual  │
│                                    ← skip3-8 │
│                                               │
│  SubbandAttention: frequency band weighting   │
└───────────────────────────────────────────────┘
    │
    ▼
CausalDPGRNN × 2
    │   pre:  Linear(32 → 32)
    │   intra: BiGRU×2 layers (freq axis, 55 steps)  ← bidirectional OK
    │   post: Linear(64 → 32)
    │   inter: UniGRU×2 layers (time axis, T steps)  ← must be unidirectional!
    │   post2: Linear(32 → 32)
    │   + LayerNorm + learnable residual scaling (α, β)
    │
    ▼
┌─ CausalDecoder ───────────────────────────────┐
│                                               │
│  CausalGTConvLite × 6 (dilation: 2,4,8,4,2,1)│
│      + skip connections (reverse order)       │
│                                               │
│  Fuse: Conv2d(64→32, k=1×1) + skip2          │
│  DSDeconv (32→32ch): 55 → 110                │
│  DSDeconv (32→2ch):  110 → 219  + skip1      │
│                                               │
└───────────────────────────────────────────────┘
    │
    ▼
out_conv: Conv2d(2 → 2, k=1×1)
    │
    ▼
ERB_48k.bs(): 219 → 513 bands
    │
    ▼
CRM mask: out = spec * mask (complex multiplication)
    │
    ▼
Output (B, 513, T, 2)

Causal Module Details

CausalGTConvLite vs GTConvLite

GTConvLite (offline):
    padding = (dilation, 1)  # symmetric padding on time axis
    can see dilation frames in both past and future

CausalGTConvLite (causal):
    pad_t = (kernel-1) * dilation = 4 * dilation
    F.pad(x, (0, 0, pad_t, 0))  # pad only on the left (past)
    can only see past frames

CausalTRA vs TRALite

TRALite (offline):
    Conv1d(ch, ch, k=5, padding=2)  # symmetric padding
    can see 2 frames in both directions

CausalTRA (causal):
    Conv1d(ch, ch, k=5, padding=0)
    F.pad(x, (4, 0))  # pad 4 frames on the left
    can only see past 4 frames

CausalDPGRNN vs DPGRNN

DPGRNN (offline):
    intra: BiGRU (freq axis)  ← no causality impact
    inter: BiGRU (time axis)  ← non-causal!

CausalDPGRNN (causal):
    intra: BiGRU (freq axis)  ← kept bidirectional
    inter: UniGRU (time axis) ← changed to unidirectional

Streaming State

The following states must be maintained for streaming inference:

• GTConv cache (12 layers, different dilation requires different buffer lengths)
• TRA history frames (12 layers, 4 frames each)
• GRU hidden states (2×DPGRNN × 2 inter layers)
• Skip connection cache (8 groups)

Directory Structure

v3_causal_stream/
├── checkpoints/
│   └── best_model_epoch35_score2.983.tar
├── configs/
│   └── cfg_causal_v2_48k.yaml
├── models/
│   ├── gtcrn_light_v3_48k_causal_v2.py    # causal model
│   └── gtcrn_light_v3_48k_causal_train.py # deprecated
├── scripts/
│   ├── inference_causal_stream.py         # batch inference
│   └── inference_stream_realtime.py       # streaming inference
├── test_samples/
└── C_Stream/                              # C implementation

Inference

Python Batch Inference

python scripts/inference_causal_stream.py -i input.wav -o output.wav \
    -c checkpoints/best_model_epoch35_score2.983.tar

Python Streaming Inference

python scripts/inference_stream_realtime.py -i input.wav -o output.wav \
    -c checkpoints/best_model_epoch35_score2.983.tar

C Streaming Inference (for deployment)

cd C_Stream && make
./build/test_audio -w weights/gtcrn_causal_v2.bin -i input.wav -o output.wav

C Implementation Performance

• Per-frame processing: ~2.1ms
• RTF: 0.21 (4.7× headroom remaining)
• Accuracy difference vs Python: ~0.01

Test Results (DNSMOS)

File	Original	Python Batch	Python Stream	C Stream
1.wav	2.21	2.85	2.69	2.69
2.wav	2.41	3.66	3.50	3.50
3.wav	2.93	3.82	3.85	3.85
4.wav	2.65	3.78	3.79	3.78
Average	2.57	3.57	3.53	3.53

The slight drop in streaming vs batch is expected, since streaming cannot look at future frames.

Training Config

width_mult: 2.0
use_two_dpgrnn: true
dpgrnn_layers: 2
loss: TransientAwareLoss (weight=1.0)
lr: 8e-4 → 5e-6
batch_size: 6
augment: random_gain [0.5, 2.0]

Streaming State & Memory Usage

Model Weights (static, CH=32)

Item	FP32	INT8 (post-QAT)
weights/biases	566.9 KB	141.7 KB
ERB filters W_bm/W_bs	128.1 KB	128.1 KB (kept FP32)
BN running stats	~6 KB	~6 KB (kept FP32)
Total	~701 KB	~276 KB

Runtime State Cache (gtcrn_state_t, CH=32)

GTConv causal cache size is determined by CH × CAUSAL_PAD_D8 × DP_W, where CAUSAL_PAD_D8=32 (history frames for max dilation=8) and DP_W=55 (frequency dimension).

Item	Size
GTConv causal cache (6 layers each in encoder/decoder)	2,703 KB
TRA time cache (12 layers × 4 frames × CH)	6.0 KB
DPGRNN inter hidden states	27.5 KB
Skip connection buffers	62.0 KB
STFT buffer	12.0 KB
Workspace	~45 KB
STFT processor	~13 KB
Total runtime	~2,869 KB

Total memory: ~3,570 KB (3.49 MB), runtime accounts for 80%, with GTConv causal cache being the dominant factor (94% of runtime).

Related Versions

• v1: Baseline
• v2: Transient optimization
• v3 (this version): Causal streaming