sonobase

自监督音频表示学习工具包：掩码谱图重建（MAE）与对比学习（NT-Xent）预训练， 面向声音事件与音频特征表示，可作为视听（audio-visual）研究的音频基座。

这是我在硕士阶段做视听表示学习时，为了给「音频侧」搭一个干净、可复用的预训练底座 而整理出来的早期基础工作。目标不是刷 SOTA，而是把特征提取、掩码、增强、损失、 骨干网络这几块拆成互相独立、便于阅读与替换的小模块。

设计取向

• NumPy 优先的核心。 谱图特征、掩码策略、数据增强、损失函数的参考实现全部只依赖 NumPy，可以在没有 GPU、甚至没有安装 PyTorch 的环境里直接跑通与单测。
• PyTorch 只在需要时出现。 骨干网络、MAE / 对比模型、训练循环放在 sonobase.models / sonobase.train 下，作为可选的 torch extra；import sonobase 本身不会引入 torch。
• 两条自监督路线共用同一个编码器。 SpectrogramTransformer 既是 MAE 的 encoder， 也是对比学习的 backbone，预训练完直接拿去做下游线性探针（linear probe）。

安装

# 仅特征 / 掩码 / 增强（纯 NumPy）
pip install -e .

# 需要神经网络与训练时，附带可选的 torch
pip install -e ".[torch]"

要求 Python 3.10+。

快速上手

1. 提取对数梅尔谱

import numpy as np
from sonobase import LogMelSpectrogram

extractor = LogMelSpectrogram(n_mels=64)          # 16 kHz / n_fft=400 / hop=160
signal = np.random.randn(16000).astype("float32")  # 1 秒音频
feats = extractor(signal)                           # -> (64, T)

2. 掩码谱图（MAE 的核心操作）

from sonobase.masking import build_masker
from sonobase.config import MaskConfig

masker = build_masker(MaskConfig(strategy="random", mask_ratio=0.75))
result = masker(feats[:, :96])                      # 谱图切成 16x16 patch
print(result.grid, result.num_masked, result.mask_ratio)

3. 掩码谱图自编码器预训练

import torch
from sonobase.config import TrainConfig
from sonobase.models import MaskedSpectrogramAutoencoder
from sonobase.train import Trainer, mae_loss

model = MaskedSpectrogramAutoencoder(mask_ratio=0.75)
trainer = Trainer(model, TrainConfig(max_steps=100), mae_loss)
trainer.fit(batch for batch in my_batches)          # batch: (B, n_mels, T)

4. 对比学习预训练

from sonobase.models import ContrastiveAudioModel
from sonobase.train import Trainer, contrastive_loss

model = ContrastiveAudioModel(temperature=0.2)
trainer = Trainer(model, TrainConfig(max_steps=100), contrastive_loss)
trainer.fit((view1, view2) for view1, view2 in my_paired_batches)

命令行

sonobase info                       # 打印版本与默认配置
sonobase features --input a.wav     # 提取并（可选）保存特征
sonobase mask-demo --strategy block # 查看某种掩码策略的统计
sonobase pretrain-mae --steps 50    # 用合成数据快速跑通 MAE 训练

更多可运行示例见 examples/。

文档

• 使用指南
• 架构总览
• 设计笔记
• API 参考

许可证

MIT © You Tao