CLAUDE.md

# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Project Overview

This is the **BEATs** (Audio Pre-Training with Acoustic Tokenizers) project, part of Microsoft's UniLM family. It implements a self-supervised audio pre-training framework based on iterative acoustic tokenization and masked audio modeling. Paper: [arXiv:2212.09058](https://arxiv.org/abs/2212.09058).

The repo also contains an active extension, **Spatial-BEATs**, which adds spatial audio understanding (direction-of-arrival, distance estimation) on top of the frozen BEATs encoder for First Order Ambisonics (FOA) data.

## Key Dependencies

- PyTorch, torchaudio (for fbank feature extraction via `torchaudio.compliance.kaldi`)
- `einops` (used by quantizer for codebook k-means init)
- Training uses `torchrun` for distributed data parallel

## Training Commands

### Spatial-BEATs (three-stage mono-AST on ov1 FOA data)
```bash
# All knobs overridable via env vars: GPUS, BATCH_SIZE, NUM_WORKERS, etc.
./run_ov1_ast_three_stage.sh
```
Stages: (1) class warmup with frozen BEATs, (2) spatial-first, (3) balanced classification + spatial.

### Pre-trunk AST experiment (two-stage)
```bash
./run_ov1_pretrunk_ast_experiment.sh
```
Stages: (1) class-only warmup with task tokens inside BEATs trunk, (2) spatial CE finetune.

### Single training run
```bash
torchrun --nproc_per_node=4 train_spatial_beats.py \
  --preset <preset_name> \
  --output-dir <output_dir> \
  --batch-size 8 --num-workers 4 --num-epochs 12
```
Available presets are defined via `make_*_config()` factories in `train_spatial_beats.py` and listed in `spatial_beats_ov123_stage1_config.py`.

## Architecture

### Original BEATs (inference-only weights)

```
Raw waveform (16kHz)
  → fbank (128 mel bins, frame_length=25ms, frame_shift=10ms)
  → normalize with fixed mean/std
  → Conv2d patch embedding
  → LayerNorm → optional Linear projection
  → TransformerEncoder (N layers with relative position bias + GRU gating)
  → extract_features() returns [B, T, D] representations
  → (finetuned models) → Linear predictor → sigmoid → class probabilities
```

Two model classes share this backbone:
- **`BEATs`** (`BEATs.py`): audio encoder. `extract_features()` returns representations or class probs (if finetuned).
- **`Tokenizers`** (`Tokenizers.py`): same encoder + `NormEMAVectorQuantizer` head. `extract_labels()` returns discrete codebook indices.

### Spatial-BEATs extension

Builds on top of BEATs to add spatial audio capabilities:

- **`SpatialBEATs`** (`spatial_beats.py`): wraps a frozen BEATs `TransformerEncoder` with multi-channel FOA preprocessing (`SpatialBEATsPreprocessor`), a `SpatialPatchEmbedding` for the extra channels, and task-specific prediction heads.
- **`spatial_modules.py`**: contains all building blocks — `SpatialPatchEmbedding`, `SpatialDeltaPatchAdapter`, `FixedSlotReadout`, `MonoTaskTokenReadout`, `FrequencyPool`, `TemporalResampler`, and prediction heads (`SpatialPredictionHeads`, `MonoTaskPredictionHeads`, `PreTrunkASTPredictionHeads`).
- **`spatial_dataset.py`**: `SpatialDataset` loads FOA audio from JSONL manifests with per-frame source annotations (azimuth, elevation, distance, class). Uses a Qwen-2.5-Omni-aligned mel frontend (16kHz, 128 bins, hop=160).
- **`spatial_loss.py`**: multi-task loss with Hungarian-style slot matching — activity BCE, azimuth/elevation CE over binned angles, distance regression, and auxiliary source classification.

### Module dependency graph

```
modules.py          — primitives: GradMultiply, SamePad, GLU_Linear, quant_noise, activation fns
quantizer.py        — NormEMAVectorQuantizer, EmbeddingEMA (VQ-VAE codebook with EMA updates)
backbone.py         — TransformerEncoder, TransformerSentenceEncoderLayer, MultiheadAttention
BEATs.py            — BEATs model (uses backbone)
Tokenizers.py       — Tokenizers model (uses backbone + quantizer)
spatial_modules.py  — spatial building blocks (patch embeddings, readout heads, prediction heads)
spatial_beats.py    — SpatialBEATs model (uses backbone + spatial_modules)
spatial_dataset.py  — SpatialDataset + collation
spatial_loss.py     — loss computation + slot matching (uses spatial_modules output types)
train_spatial_beats.py — training loop, presets, CLI (uses spatial_beats, spatial_dataset, spatial_loss)
```

## Loading Pre-trained Checkpoints

Checkpoints are `dict` with keys `'cfg'` (config dict) and `'model'` (state dict):
```python
checkpoint = torch.load('model.pt')
cfg = BEATsConfig(checkpoint['cfg'])
model = BEATs(cfg)
model.load_state_dict(checkpoint['model'])
```
Same pattern for `Tokenizers` with `TokenizersConfig`.

## Audio Input Contract

- All models expect **16kHz mono** waveforms
- `preprocess()` converts to 128-bin fbank features normalized with fixed mean=15.41663, std=6.55582
- Padding masks are `bool` tensors where `True` = padded position
- Spatial-BEATs uses 4-channel FOA input instead of mono

我希望在原始BEATs的基础上更改模型的框架，让模型有FOA音频的理解能力，能够在声源分类之外拥有识别位置的能力，这样的encoder作为我未来输入给LLM的例子。我之前自己尝试了一些做法，不过class分类不是很收敛，空间指标比如dis,ele,azimuth的loss几乎不收敛，我感觉我的方法太过于ML了，没有充分的利用DL的能力，或许应该一定程度上相信attention的能力来学习。我认为应该像BAT一样，你看这个目录下面的Spatial-AST的训练是从AudioMAE的训练开始的，我觉得确实应该学习他的设计来类似的训练我的Spatial-BEATs，我设计了实验run_ov1_pretrunk_ast_experiment.sh来验证，现在有了初步的结果，但是看的出来，还不是很收敛，预期结果和我想的完全不一样，我到底应该怎么办呢？还有疑问是BEATS是用audioset训练的，我现在的ov1数据干声来源于FSD50K，这是不是首先会影响分类任务，我是不是应该先在分类任务上finetune到一定的程度之后再考虑空间呢