finetune_spk-sortformer

G-STAR Sortformer 微调流水线

多说话人日志化-AST（自动语音转录）微调流水线，适用于 G-STAR 流式 Sortformer 模型。通过 FastMSS 生成合成多说话人对话数据，并使用 split-head 差异化学习率训练支持任意说话人数量扩展（4→N）的日志化 ASR 模型。

核心功能

• 多说话人数据模拟 — 基于 FastMSS 的会议对话生成，包含 HMM 话轮切换、RIR 混响、噪声增强和三角数加权说话人采样
• Sortformer 微调 — 流式 Sortformer 训练，支持 Hungarian PIT 损失、split-head 架构和 bf16 精度
• 任意 N 说话人扩展 — 基于 SVD 正交初始化，将预训练的 4 说话人检查点扩展到 N 说话人，无需从头训练
• 流式推理 — 基于 FIFO 说话人缓存的逐块日志化，支持 ATS（到达时间排序）标签分配
• 数据分析 — 生成数据的统计分析（说话人分布、时长分布及 CPD 检测）
• 新数据源集成 — 支持通过 G-STAR 对齐流水线扩展到无预对齐的音视频+文本数据源

目录结构

finetune_pipeline/
├── train.sh                              # Sortformer 训练入口
├── infer.sh                              # 流式推理入口
├── run_train.sh                          # 训练示例脚本
├── run_infer.sh                          # 推理示例脚本
├── path.sh                               # Python 环境配置
├── requirements.txt                      # 补充 Python 依赖（NeMo/torch 之外）
│
├── src/
│   ├── finetune_pipeline/
│   │   ├── bin/
│   │   │   ├── streaming_sortformer_diar_train.py   # 训练入口脚本
│   │   │   └── e2e_diarize_speech.py                # 推理入口脚本
│   │   ├── scripts/
│   │   │   ├── ckpt2nemo.py               # Lightning ckpt → NeMo .nemo 格式
│   │   │   ├── extend_output_layer.py     # SVD 正交初始化 N 说话人扩展
│   │   │   ├── merge_split_head.py        # 微调后合并 split head → 统一层
│   │   │   ├── convert_infer_log_to_csv.py # 推理日志 → CSV 格式
│   │   │   ├── split_manifest_by_spk.py   # 按说话人数量分割 manifest
│   │   │   └── parse_options.sh           # Bash 命令行参数解析器
│   │   └── conf/
│   │       ├── streaming_sortformer_diarizer_4spk-v2.yaml
│   │       ├── streaming_sortformer_diarizer_8spk-v1-finetune.yaml
│   │       └── streaming_sortformer_diarizer_10spk-v1.yaml
│   └── third_party/
│       ├── FastMSS/          # FastMSS 分支（子模块）— 数据模拟引擎
│       └── nemo/             # 裁剪版 NeMo — Sortformer 模型及工具
│
├── data/
│   ├── fastmss/              # FastMSS 数据模拟配置及脚本
│   │   ├── simulate.sh       # FastMSS 单次运行入口
│   │   └── scripts/          # CutSet 生成、分析等工具
│   ├── shared/               # 共享工具脚本
│   └── conf/                 # 模拟 Hydra 配置
│
├── docs/                     # 项目文档（见下方索引）
├── checkpoints/              # 预训练模型检查点
├── exp/                      # 训练实验输出
└── dev/                      # 开发/调试脚本

快速开始

前置条件

从零搭建环境（与参考部署一致）：

# 1. 创建 Python 3.10 环境
conda create -n nemo python=3.10 -y
conda activate nemo

# 2. 安装 PyTorch 2.9.1 (CUDA 12.6)
pip install torch==2.9.1 torchcodec==0.9.1 torchaudio==2.9.1 \
  --index-url https://download.pytorch.org/whl/cu126

# 3. 系统依赖 + NeMo toolkit（装依赖用；运行时由 path.sh 优先使用 vendored NeMo）
apt-get update && apt-get install -y libsndfile1 ffmpeg
# 无 sudo 时可用: conda install ffmpeg -c conda-forge
pip install Cython packaging
pip install "nemo_toolkit[asr] @ git+https://github.com/NVIDIA/NeMo.git@main"

# 4. 安装本仓库补充依赖
pip install -r requirements.txt

每次运行训练/推理/仿真前先激活环境：

. ./path.sh

path.sh 会将 src/third_party 置于 PYTHONPATH 最前并 conda activate nemo，因此 import nemo 实际加载的是仓库内 vendored/patched NeMo，与上面 pip 安装的 nemo_toolkit 依赖共存。若 conda 路径不同，请修改 path.sh 中的激活方式。

1. 生成训练数据 (FastMSS)

# 完整的 train/dev/test 模拟（6 说话人，总计 750 小时）
bash data/fastmss/run_6spk_v2.sh

# 单次运行（自定义参数）
bash data/fastmss/simulate.sh \
    --output-dir data/dump_fastmss/my_run \
    --config-path extend_sortformer/6spk \
    --config-name v2 \
    --profile train \
    --spk-weights "[1,3,6,10,15,21]" \
    --n-meetings 30000 \
    --seed 42

2. 训练模型

# 使用预配置的 4 说话人训练
bash run_train.sh \
    --train_manifest_path data/train_manifest.json \
    --dev_manifest_path data/dev_manifest.json \
    --exp_name 4spk_v1

# 或使用原始 train.sh
bash train.sh \
    --train_manifest_path data/train_manifest.json \
    --dev_manifest_path data/dev_manifest.json \
    --exp_name 4spk_v1

3. 推理

# 单数据集推理
bash run_infer.sh \
    --model_path checkpoints/model.nemo \
    --dataset_manifest data/test_manifest.jsonl \
    --exp_dir exp/infer/test

# 批量推理多个数据集
bash run_infer.sh \
    --model_path checkpoints/model.nemo \
    --dataset_base_dir data/global_testset \
    --datasets "ami fisher mlc" \
    --exp_dir exp/infer/batch_test

支持任意说话人数量的微调方法

概述

Sortformer 原始设计支持固定数量的说话人（如 4 说话人）。为了扩展到更多说话人（N>4），本项目提供了基于 split-head 架构的扩展方法：

1. 扩展输出层：将原始的 Linear(192, 4) 分解为 Linear(192, 4) + Linear(192, N-4)
2. 差异化学习率：基础说话人使用较低学习率（1e-5），新说话人使用较高学习率（1e-4）
3. 训练后合并：将 split head 合并为统一的 Linear(192, N) 便于后续推理或继续微调

完整流程

步骤 1：扩展输出层（SVD 正交初始化）

# 将 4 说话人模型扩展到 8 说话人
python src/finetune_pipeline/scripts/extend_output_layer.py \
    --src checkpoints/pretrained/sortformer_4spk.nemo \
    --dst-spk 8 \
    --out checkpoints/8spk_extended.nemo

原理说明：

• 对原始权重矩阵 W 进行 SVD 分解：U, S, Vh = torch.linalg.svd(W, full_matrices=True)
• 新说话人的权重行取自 Vh[4:]，与原有行空间正交
• 按原矩阵平均范数进行归一化，保持数值稳定性

步骤 2：使用差异化学习率微调

# 使用扩展后的模型进行微调
bash run_train.sh \
    --train_manifest_path data/train_8spk.json \
    --dev_manifest_path data/dev_8spk.json \
    --config_name streaming_sortformer_diarizer_8spk-v1-finetune.yaml \
    --init_nemo_path checkpoints/8spk_extended.nemo \
    --exp_name 8spk_finetune

配置要求： 在配置文件中设置：

model:
  sortformer_modules:
    num_spks: 8
    n_base_spks: 4  # 启用 split-head 架构
  lr: 1e-5          # 基础说话人学习率
  optim_new_lr: 1e-4  # 新说话人学习率

步骤 3：合并 Split Head（可选但推荐）

# 微调完成后合并 split head
python src/finetune_pipeline/scripts/merge_split_head.py \
    --src exp/streaming_sortformer_diar_train/8spk_finetune/checkpoints/model--val_loss=xxx.nemo \
    --out checkpoints/8spk_merged.nemo

何时需要合并：

• ✅ 需要继续微调到更多说话人（如 8→10）
• ✅ 需要简化推理模型结构
• ✅ 需要发布最终模型
• ❌ 仅用于单次推理（可直接使用 split-head 模型）

示例：4 → 8 → 10 说话人扩展

# 第一阶段：4 → 8 说话人
python src/finetune_pipeline/scripts/extend_output_layer.py \
    --src checkpoints/4spk.nemo --dst-spk 8 --out checkpoints/8spk_extended.nemo

bash run_train.sh \
    --config_name streaming_sortformer_diarizer_8spk-v1-finetune.yaml \
    --init_nemo_path checkpoints/8spk_extended.nemo \
    --exp_name 4to8_finetune

# 合并 8 说话人模型
python src/finetune_pipeline/scripts/merge_split_head.py \
    --src exp/.../checkpoints/model.nemo --out checkpoints/8spk_merged.nemo

# 第二阶段：8 → 10 说话人
python src/finetune_pipeline/scripts/extend_output_layer.py \
    --src checkpoints/8spk_merged.nemo --dst-spk 10 --out checkpoints/10spk_extended.nemo

bash run_train.sh \
    --config_name streaming_sortformer_diarizer_10spk-v1.yaml \
    --init_nemo_path checkpoints/10spk_extended.nemo \
    --exp_name 8to10_finetune

注意事项

1. 数据要求：微调数据必须包含目标说话人数量的样本（如 8 说话人训练需要 6-8 说话人的数据）
2. 学习率设置：新说话人学习率应比基础说话人高约 10 倍
3. 训练稳定性：建议先用小学习率预热，再逐步增大
4. 内存需求：split-head 模型比统一模型占用稍多内存

文档索引

依赖项

• Python: PyTorch, NeMo (裁剪版), Lhotse, Pyroomacoustics, Omegaconf, SoundFile
• 外部工具: Qwen3-ForcedAligner（用于基于对齐的数据源扩展）
• 数据: LibriSpeech, WHAM 噪声数据
• 硬件: A800-80G / A100 GPU（已测试 bf16 训练）

许可证

本项目包含来自 NVIDIA NeMo (Apache 2.0) 和 FastMSS (MIT) 的裁剪代码。详细信息请参阅各源文件头部。