RE-USE(Rethinking Targets, Architectures and Data Quality for Universal Speech Enhancement)是 NVIDIA 提出的通用语音增强模型,支持多采样率(8/16/22.05/24/32/44.1/48 kHz)、多退化类型(噪声、混响、削波、带宽限制等),具有语言无关性。
本项目将 RE-USE 模型适配到 华为昇腾 Ascend NPU 平台,使用 torch_npu 替代原有的 CUDA/Mamba 计算栈进行推理。
| 模块 | 原实现 | 昇腾适配方案 |
|---|---|---|
| Mamba 扫描 | CUDA selective_scan_fn | PyTorch 循环实现 (models/mamba_npu.py) |
| STFT/iSTFT | PyTorch GPU | STFT on NPU + iSTFT CPU 回退 |
| 权重格式 | PyTorch .pth | Safetensors (model.safetensors) |
| 音频 I/O | torchaudio | scipy.io.wavfile |
RE-USE_ascend/)| 组件 | 版本 |
|---|---|
torch | 2.6.0 |
torch_npu | 2.6.0.post1 |
CANN | 8.5.1 |
| NPU 设备 | Ascend910B4(29.49 GB) |
| Python | 3.11.14 |
librosa | 0.10.2 |
scipy | 1.14.1 |
safetensors | 0.5.3 |
pip install -i https://repo.huaweicloud.com/repository/pypi/simple \
torch torch_npu safetensors scipy numpy librosa soundfile resampy# 从 ModelScope 下载(推荐,国内网络)
pip install modelscope
python -c "from modelscope import snapshot_download; snapshot_download('nvidia/RE-USE', local_dir='./RE-USE')"
# 或从 HuggingFace 镜像
# pip install huggingface-hub
# huggingface-cli download --resume-download nvidia/RE-USE --local-dir ./RE-USE --local-dir-use-symlinks Falsecd /path/to/RE-USE_ascend
# 直接运行
python inference.py \
--input ./noisy_audio/mic_test2.wav \
--output ./output/enhanced_output.wav \
--checkpoint ./model.safetensors \
--config ./config.json
# 或使用脚本
bash inference.sh| 参数 | 默认值 | 说明 |
|---|---|---|
--input | noisy_audio/mic_test2.wav | 输入含噪音频文件 |
--output | output/enhanced_output.wav | 输出增强音频文件 |
--config | config.json | 模型配置文件(JSON) |
--checkpoint | model.safetensors | 模型权重文件(safetensors) |
# 基础推理验证
python inference.py --input noisy_audio/mic_test2.wav --output output/enhanced_output.wav
# 预期输出
# [INFO] Using device: npu:0
# [INFO] NPU device name: Ascend910B4
# [INFO] Model parameters: 9,607,747 (9.61M)
# [INFO] Inference time: 3.300s (RTF=1.55x)
# [INFO] Enhanced audio saved to: output/enhanced_output.wav验证结果:
npu:0 设备可用 ✅测试条件:noisy_audio/mic_test2.wav(5.10s, 44.1kHz → 重采样至 8kHz),5 次推理取平均。
| 指标 | 数值 |
|---|---|
| 输入音频时长 | 5.10 s |
| 推理耗时(平均) | 3.300 s |
| 推理耗时(中位数) | 3.269 s |
| RTF(实时倍率) | 1.55x |
| NPU 峰值内存 | 1792 MB |
| 模型参数量 | 9.61M |
| 输出文件大小 | 159.4 KB |
说明:RTF 1.55x 意味着处理 1 秒音频约需 0.65 秒(慢于实时)。瓶颈在 Mamba NPU 的 PyTorch 循环实现。使用
torch.compile或昇腾亲和算子可进一步提升。
使用 eval_accuracy.py 进行有参考精度评测(输入含噪 vs 输出增强),5.10s 音频在 8kHz 对齐比较。
| 指标 | 数值 |
|---|---|
| 输入 RMS | 0.040 |
| 输出 RMS | 0.032 |
| 能量比(out/in) | -1.91 dB |
| SI-SNR | 2.79 dB |
| 波形相关系数 | 0.8094 |
| 输入峰值 | 0.5388 |
| 输出峰值 | 0.6572 |
| 频谱活动帧(输入) | 0.9% |
| 频谱活动帧(输出) | 0.8% |
| 频谱距离 (MSE dB) | 434.56 |
说明:输入为含噪语音(非纯净参考),因此 SI-SNR 2.79 dB 和相关系数 0.81 表示输出在保持语音结构的同时进行了降噪处理。
核心验证:同一段音频(1.0s, 8kHz, 8000 采样点)分别在 CPU(PyTorch) 和 NPU(torch_npu) 上推理,逐采样点计算误差。
| 指标 | 数值 |
|---|---|
| 对比总采样点 | 8000 |
| 信号活动区占比 | 23.6% |
| 信号活动区平均相对误差 | 0.395% ✅ |
| 余弦相似度 | 0.999996 |
| 误差 SNR | 51.91 dB |
| 全局 MAE | 5.27e-05 |
| 全局 RMSE | 1.64e-04 |
| 信号区 P95 绝对误差 | 7.98e-04 |
结论:✅ 信号活动区平均相对误差 0.395% < 1%,NPU 精度在 1% 以内。
余弦相似度 0.999996 表明波形几乎完全一致,51.9 dB 的误差 SNR 表明差异处于浮点运算舍入误差量级,不影响听觉质量。
差异来源:NPU 与 CPU 的浮点运算顺序不完全相同(
torch_npu的某些算子实现与torchCPU 路径存在汇编级差异),导致最后几位 bit 的舍入不同。这是硬件平台差异的正常现象。
原始 Mamba 使用 CUDA 自定义算子 selective_scan,在昇腾 NPU 上不可用。适配方案:
models/mamba_npu.py — 纯 PyTorch 循环实现的双向 Mamba 扫描torch.cumsum 替代 CUDA 并行扫描昇腾 NPU 的 UnfoldGrad 算子存在非连续张量问题,iSTFT 当前回退到 CPU 执行:
if com.device.type == 'npu':
com = com.cpu()
hann_window = hann_window.cpu()
wav = torch.istft(com, n_fft, ...)避免使用 torchaudio(与 torchcodec 有版本冲突),改用 scipy.io.wavfile。
输入 WAV → scipy 读取 → int16→float32 → librosa 重采样至 8kHz
→ STFT (NPU) → Mamba 增强 (NPU) → iSTFT (CPU)
→ pad/trim → scipy 写入 → 输出 WAV.pth 需自行转换RE-USE_ascend/ # 昇腾适配包
├── inference.py # 推理入口脚本
├── inference.sh # 推理启动脚本
├── config.json # 模型配置
├── model.safetensors # 模型权重
├── README.md # 本文档
├── log.txt # 精度+性能评测日志
├── eval_accuracy.py # 精度评测脚本
├── eval_perf.py # 性能评测脚本
├── run_all_eval.sh # 一键评测脚本
├── models/
│ ├── generator_SEMamba_time_d4.py # 模型主架构
│ ├── mamba_block2_SEMamba.py # Mamba Block
│ ├── mamba_npu.py # Mamba NPU 适配
│ ├── stfts.py # STFT/iSTFT 适配
│ └── codec_module_time_d4.py # Codec 模块
├── utils/
│ └── util.py # 工具函数
└── output/
└── enhanced_output.wav # 增强输出示例