z
zkx_/RE-USE
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

RE-USE on Ascend NPU — 通用语音增强昇腾适配

1. 简介

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_fnPyTorch 循环实现 (models/mamba_npu.py)
STFT/iSTFTPyTorch GPUSTFT on NPU + iSTFT CPU 回退
权重格式PyTorch .pthSafetensors (model.safetensors)
音频 I/Otorchaudioscipy.io.wavfile

获取地址

  • 原始权重(HuggingFace):https://huggingface.co/nvidia/RE-USE
  • 原始权重(ModelScope):https://modelscope.cn/models/nvidia/RE-USE
  • 昇腾适配代码:本仓库(RE-USE_ascend/)

参考文档

  • RE-USE 论文
  • NVIDIA RE-USE HuggingFace
  • torch_npu 文档

2. 验证环境

组件版本
torch2.6.0
torch_npu2.6.0.post1
CANN8.5.1
NPU 设备Ascend910B4(29.49 GB)
Python3.11.14
librosa0.10.2
scipy1.14.1
safetensors0.5.3

3. 环境准备

3.1 安装依赖

pip install -i https://repo.huaweicloud.com/repository/pypi/simple \
  torch torch_npu safetensors scipy numpy librosa soundfile resampy

3.2 下载模型权重

# 从 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 False

4. 模型推理

4.1 快速启动

cd /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

4.2 命令行参数

参数默认值说明
--inputnoisy_audio/mic_test2.wav输入含噪音频文件
--outputoutput/enhanced_output.wav输出增强音频文件
--configconfig.json模型配置文件(JSON)
--checkpointmodel.safetensors模型权重文件(safetensors)

4.3 输入输出说明

  • 输入:任意采样率(8–48 kHz)的单声道 WAV 文件
  • 输出:与模型原生采样率(8 kHz)一致的增强 WAV 文件
  • 自动处理:int16 → float32 转换、重采样至 8 kHz

5. Smoke 验证

# 基础推理验证
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 设备可用 ✅
  • 模型权重加载成功 ✅
  • 推理完成无报错 ✅
  • 输出 WAV 文件可播放 ✅

6. 性能参考

测试条件: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 或昇腾亲和算子可进一步提升。

7. 精度评测

7.1 有参考精度评测(含噪输入 vs 增强输出)

使用 eval_accuracy.py 进行有参考精度评测(输入含噪 vs 输出增强),5.10s 音频在 8kHz 对齐比较。

指标数值
输入 RMS0.040
输出 RMS0.032
能量比(out/in)-1.91 dB
SI-SNR2.79 dB
波形相关系数0.8094
输入峰值0.5388
输出峰值0.6572
频谱活动帧(输入)0.9%
频谱活动帧(输出)0.8%
频谱距离 (MSE dB)434.56

说明:输入为含噪语音(非纯净参考),因此 SI-SNR 2.79 dB 和相关系数 0.81 表示输出在保持语音结构的同时进行了降噪处理。

7.2 NPU vs CPU 逐采样点精度对比

核心验证:同一段音频(1.0s, 8kHz, 8000 采样点)分别在 CPU(PyTorch) 和 NPU(torch_npu) 上推理,逐采样点计算误差。

指标数值
对比总采样点8000
信号活动区占比23.6%
信号活动区平均相对误差0.395% ✅
余弦相似度0.999996
误差 SNR51.91 dB
全局 MAE5.27e-05
全局 RMSE1.64e-04
信号区 P95 绝对误差7.98e-04

结论:✅ 信号活动区平均相对误差 0.395% < 1%,NPU 精度在 1% 以内。

余弦相似度 0.999996 表明波形几乎完全一致,51.9 dB 的误差 SNR 表明差异处于浮点运算舍入误差量级,不影响听觉质量。

差异来源:NPU 与 CPU 的浮点运算顺序不完全相同(torch_npu 的某些算子实现与 torch CPU 路径存在汇编级差异),导致最后几位 bit 的舍入不同。这是硬件平台差异的正常现象。

8. 适配注意事项

8.1 Mamba NPU 循环实现

原始 Mamba 使用 CUDA 自定义算子 selective_scan,在昇腾 NPU 上不可用。适配方案:

  • models/mamba_npu.py — 纯 PyTorch 循环实现的双向 Mamba 扫描
  • torch.cumsum 替代 CUDA 并行扫描
  • 支持双向扫描(前向 + 后向)

8.2 iSTFT CPU 回退

昇腾 NPU 的 UnfoldGrad 算子存在非连续张量问题,iSTFT 当前回退到 CPU 执行:

if com.device.type == 'npu':
    com = com.cpu()
    hann_window = hann_window.cpu()
wav = torch.istft(com, n_fft, ...)

8.3 音频 I/O

避免使用 torchaudio(与 torchcodec 有版本冲突),改用 scipy.io.wavfile。

8.4 推理流程

输入 WAV → scipy 读取 → int16→float32 → librosa 重采样至 8kHz
  → STFT (NPU) → Mamba 增强 (NPU) → iSTFT (CPU)
  → pad/trim → scipy 写入 → 输出 WAV

8.5 已知限制

  • 多卡推理:当前仅支持单 NPU 卡
  • 带宽扩展(BWE):未适配
  • 批量推理:未实现
  • 权重格式:仅支持 safetensors,如需 .pth 需自行转换

9. 文件结构

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            # 增强输出示例

10. 参考文献

  • Rethinking Training Targets, Architectures and Data Quality for Universal Speech Enhancement, arXiv:2603.02641, 2025
  • SEMamba: Mamba-based Speech Enhancement, https://github.com/RoyChao19477/SEMamba
  • torch_npu, https://www.hiascend.com/document/detail/zh/Pytorch/60RC1/
  • NVIDIA RE-USE, https://huggingface.co/nvidia/RE-USE