本文档记录 RE-USE(Real-time Enhancement with U-shaped Speech Enhancement)在昇腾 NPU(Ascend 910B3)环境的适配部署与精度验证结果。
RE-USE 是 NVIDIA 提出的基于 Mamba 架构的实时语音增强模型。
相关地址:
| 组件 | 版本 |
|---|---|
| Python | 3.11.x |
| PyTorch | 2.10.0+cpu |
| torch_npu | 2.10.0 |
| CANN | 8.5.1 |
| NPU 硬件 | Ascend 910B3 (8 卡) |
| 项目 | 值 |
|---|---|
| 模型架构 | DenseEncoder + 30x TFMambaBlock + MagDecoder + PhaseDecoder |
| 参数量 | ~9.64M |
| 采样率 | 8kHz |
| 权重格式 | safetensors |
| 框架 | PyTorch |
| 预训练数据 | AudioSet + DNS |
| 许可证 | NVIDIA |
RE-USE 模型核心依赖 mamba_ssm.Mamba,该模块使用自定义 CUDA 内核实现选择性扫描算法,无法在昇腾 NPU 上直接运行。
本适配使用纯 PyTorch 实现的 Mamba 块(models/mamba_npu.py)替换 CUDA 内核,仅使用标准 PyTorch ops,确保在 NPU 上可运行。经验证,NPU 与 CPU 的复数综合误差 < 0.15%,满足精度要求。
conda create -n re-use python=3.11 -y
conda activate re-use
pip install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 --index-url https://repo.huaweicloud.com/repository/pypi/simple/
pip install torch_npu==2.10.0 --index-url https://repo.huaweicloud.com/repository/pypi/simple/
pip install soundfile safetensors pyyaml einops# NPU 推理
python3 inference.py --model_path /path/to/RE-USE --input_audio /path/to/noisy.wav --output_audio enhanced.wav
# CPU 推理
python3 inference.py --model_path /path/to/RE-USE --input_audio /path/to/noisy.wav --device cpu
# 评测
python3 benchmark.py --model_path /path/to/RE-USE| 输出张量 | 向量级相对误差 | 余弦相似度 |
|---|---|---|
| magnitude | 0.052% | 1.000000 |
| phase | 0.161% | — |
| complex (mag+phase combined) | 0.146% | 0.999999 |
| 指标 | 实测值 | 阈值 | 状态 |
|---|---|---|---|
| 复数综合误差(主指标) | 0.146% | < 1% | ✅ PASS |
| 幅度误差 | 0.052% | < 1% | ✅ PASS |
| 相位误差(非循环比较) | 0.161% | < 1% | ✅ PASS |
| 操作 | 耗时 |
|---|---|
| NPU 推理时间(3s 音频) | 1.83s |
| 实时因子(RTF) | 0.61 |
huggingface_hub 下载权重,本适配改为本地加载。