PS4 is a
Target Speaker Extraction (TSE) model that jointly optimizes speech separation quality and ASR transcription accuracy through proxy supervision. It is fine-tuned from a pretrained BSRNN + ECAPA-TDNN backbone on the
REAL-PS4 dataset — a multi-domain real-recording corpus covering Chinese and English meeting scenarios.
Trained on
REAL-PS4, which aggregates four real-recording meeting datasets:
Use the included
inference.py — a self-contained script with absolutely no external dependencies beyond
torch,
torchaudio, and
numpy:
1# Install dependencies
2pip install torch torchaudio numpy
3
4# Single file extraction
5python inference.py \
6 --checkpoint checkpoint_epoch037.pt \
7 --mix mix.wav \
8 --enroll target_speaker.wav \
9 --output result.wav
10
11# Use GPU
12python inference.py \
13 --checkpoint checkpoint_epoch037.pt \
14 --mix mix.wav \
15 --enroll target.wav \
16 --output result.wav \
17 --device cuda
18
19# Batch mode (process all .wav files in a directory)
20python inference.py \
21 --checkpoint checkpoint_epoch037.pt \
22 --mix-dir ./mixtures/ \
23 --enroll-dir ./enrollments/ \
24 --output-dir ./results/ \
25 --device cuda
26
27# List available CUDA devices
28python inference.py --list-devices
1import torch
2from inference import BSRNN, load_audio, extract_speaker, load_checkpoint, build_model
3
4# Build model and load weights
5device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6model = build_model(device)
7load_checkpoint("checkpoint_epoch037.pt", model, device)
8
9# Load audio (16 kHz mono)
10mix = load_audio("mixture.wav")
11enroll = load_audio("enrollment.wav")
12
13# Run extraction
14extracted = extract_speaker(model, mix, enroll, device)
15
16# Save result
17torchaudio.save("result.wav", extracted, 16000)
1import torch
2import torchaudio
3from inference import BSRNN
4
5# Build model with exact training config
6model = BSRNN(
7 feat_type="consistent",
8 feature_dim=128,
9 num_repeat=6,
10 spk_emb_dim=192,
11 spk_fuse_type="multiply",
12 multi_fuse=False,
13 spk_model="ECAPA_TDNN_GLOB_c512",
14 sr=16000, win=512, stride=128,
15 spk_args={"feat_dim": 80, "embed_dim": 192, "pooling_func": "ASTP"},
16 spk_model_freeze=True,
17 use_spk_transform=False,
18 joint_training=True,
19 multi_task=False,
20 spk_feat=False,
21)
22model.eval()
23
24# Load checkpoint
25ckpt = torch.load("checkpoint_epoch037.pt", map_location="cpu")
26state_dict = ckpt["model"] if "model" in ckpt else ckpt
27model.load_state_dict(state_dict, strict=False)
28
29# Run inference
30mix, sr = torchaudio.load("mixture.wav")
31enroll, sr = torchaudio.load("enrollment.wav")
32with torch.no_grad():
33 extracted, _ = model(mix, enroll)
1@misc{ning2026ps4,
2 title = {PS4: Proxy-Supervised Joint Training for Real Target Speaker Extraction},
3 author = {Wanyi Ning and Wei Zhou and Yingpeng Li and Yinshang Guo and Haitao Qian and Yiming Cheng},
4 year = {2026},
5 eprint = {2607.08111},
6 archivePrefix = {arXiv},
7 primaryClass = {cs.SD},
8 url = {https://arxiv.org/abs/2607.08111}
9}
