Nemotron-3-Diarization.AXERA

nvidia/Nemotron-3-Diarization（NVIDIA Open Model License，openmdw-1.1）端到端流式说话人日志 在 AX650N 的量化部署包：输入 16 kHz 单声道会议音频，输出 RTTM 说话人时序标签 （最多 8 说话人，10 ms 输出分辨率）。

• 模型: 8 张拆图 axmodel（U16 激活 / S8 权重），两档流式配置 × 两个 NPU 核数版本
  • offline（精度优先，输入缓冲 30.4 s）：preencode.axmodel + encoder.axmodel（NPU3 / 三核）； 双核版 preencode_npu2.axmodel + encoder_npu2.axmodel
  • low latency（实时流式 ~0.8 s）：preencode_lowlat.axmodel + encoder_lowlat.axmodel； 双核版 preencode_lowlat_npu2.axmodel + encoder_lowlat_npu2.axmodel
  • NPU2 版仅编译期调度不同：输出与 NPU3 版逐帧完全一致（maxdiff 0），单步耗时更高， 适合把 1 个 NPU 核让给其他模型并跑的场景
• 依赖: Python 仅 numpy + soundfile + axengine；C++ 为预编译 aarch64 可执行文件
• 转换源码（导出/对分/校准/Pulsar2 编译）见 GitHub: ZY-2012/Nemotron-3-Diarization.AXERA

目录

├── models/          # 4 张 axmodel + 运行时资产（learnable_sil_emb/feature_window/mel_filterbank）+ model_meta
├── python/          # Python 推理入口（nemotron_sdk + example.py）
├── bin/             # C++ 可执行文件（aarch64）：nemotron_ax650
├── samples/         # 演示音频（AMI 会议 120 s）
├── run_ax650.sh  run_cpp_ax650.sh   # 一键运行（Python / C++，--lowlat 切档）
└── requirements.txt

Python 运行

pip3 install numpy soundfile
# pyaxengine: https://github.com/AXERA-TECH/pyaxengine/releases/latest
pip3 install axengine-<version>-py3-none-any.whl

./run_ax650.sh                              # 自带样例 samples/sample_meeting.wav
./run_ax650.sh input.wav out.rttm --lowlat  # 低延迟档
./run_ax650.sh input.wav out.rttm --npu2    # 双核编译版（可与 --lowlat 组合）

或:

import sys; sys.path.insert(0, "python")
from nemotron_sdk import AxengineGraphPair, StreamingDiarizer, SortformerConfig, load_learnable_sil_emb

config = SortformerConfig()                       # offline 30.4 s 档
load_learnable_sil_emb(config, "models/learnable_sil_emb.npy")
diarizer = StreamingDiarizer(
    AxengineGraphPair("models/preencode.axmodel", "models/encoder.axmodel"),
    config,
    feature_assets="models",                      # 复用检查点内量化窗/滤波器组
)
preds = diarizer.process_wav(waveform_16k_float32, 16000)   # (T, 8) @ 10 ms
lines = diarizer.rttm_lines(preds, uri="meeting")

C++ 运行

export LD_LIBRARY_PATH=/soc/lib:${LD_LIBRARY_PATH:-}
./bin/nemotron_ax650 --preencode models/preencode.axmodel --encoder models/encoder.axmodel \
    --sil-emb models/learnable_sil_emb.npy --assets models \
    --wav input.wav --rttm out.rttm --threads 8
# 或一键：./run_cpp_ax650.sh input.wav out.rttm [--lowlat] [--npu2]

模型说明

• 输入: 16 kHz 单声道 wav；输出: RTTM（SPEAKER <uri> 1 <start> <dur> <NA> <NA> speaker_k <NA>）
• 流程: 主机 log-mel（128 mel / 25 ms 窗 / 10 ms 步长，复用检查点内量化的 hann 窗与滤波器组） → preencode（FeatureStacking，3040 mel → 380×512）→ 主机打包 seq [1,684,512] → encoder（31 层 RoPE Transformer，10 ms 高分辨率概率 [1,5472,8]） → 主机 80 ms 降采样 + streaming_update（learnable 静音嵌入 / top-k + AOSC 压缩 / FIFO 弹出） → 后处理
• 拆图契约：offline = spkcache 264 + fifo 40 + chunk 380 帧（seq 684）；low latency = 264 + 264 + 10（seq 538）； 图内无 Scatter/Where、无动态 shape；models/model_meta*.json 给出完整接口与帧率

精度与性能（AX650N 实测）

三方对比（collar=0，pyannote.metrics，本地 AMI Mix-Headset / AliMeeting 重构参考）：

数据集	NeMo FP32（参考）	主机 ONNX FP32	板端 U16/S8	ΔDER
AMI-SDM 5 场（offline）	29.82%	29.83%	29.99%（C++ 29.97%）	+0.17pp
AliMeeting 4 场（远场，offline）	17.92%	17.93%	18.23%（C++ 18.19%）	+0.31pp
AMI 2 场（low latency）	28.28%	28.26%	29.08%	+0.80pp

配置	版本	Python RTF	C++ RTF	C++ ms/步
offline（seq 684 × 31 层）	NPU3	0.0096	0.0066	177
offline	NPU2	—	0.0075	201
low latency（seq 538 × 31 层）	NPU3	0.189	0.179	85.9
low latency	NPU2	—	0.338	162

NPU2/NPU3 两版输出逐帧一致（maxdiff 0，ES2004a 实测）；NPU2 留出 1 个 NPU 核， 可用于与 ASR 等模型共享 AX650N 三核 NPU。