Skip to content

TRAINING.md

Latest commit

 

History

History
177 lines (133 loc) · 5.68 KB

TRAINING.md

File metadata and controls

177 lines (133 loc) · 5.68 KB
Raw

Training

How to train model variants on MLAAD.

Prerequisites

  • Pre-encoded dataset at data/prepared_ds_seg_enc/ (see SETUP.md → make prepare).
  • For hierarchical models: a superclass LUT CSV under data/ (see "Hierarchical setup" below).

All training scripts live in scripts/training/ and run on pre-encoded features (--pre_encoded True).

Variants

The architecture is the same Wav2Vec2 + AASIST encoder; the head and loss change per variant.

Variant Script What it is
Flat train_flat.py Single softmax head over all classes
H-Shared train_hier.py --hierarchy_type "H-Shared" Shared trunk, two heads (superclass + class), multi-task loss
H-Arch train_hier.py --hierarchy_type "H-Arch" Architecture-specific: one per-superclass head, gated by the superclass prediction

Flat — ArcFace margin

uv run scripts/training/train_flat.py \
  --path_to_dataset "data/prepared_ds_seg_enc" \
  --pre_encoded True --pre_augmented True --is_segmented True \
  --out_folder "exp/trained_models/arc_m03" \
  --use_arc_margin --easy_margin False --arc_m 0.3 \
  --weighted_sampling True

Or via makefile:

make train_flat_arc

Flat — Sub-center ArcFace

make train_flat_subarc

Equivalent to the flat ArcFace command with --use_sub-center-arc_margin --k_centers 3 --arc_m 0.3 --easy_margin True.

Hierarchical setup

Both H-Shared and H-Arch require a --superclass_lut CSV with columns Class, Superclass, Class ID, Superclass ID. Class ID must match the class IDs in data/label_assignment.txt.

The canonical full-24-class LUT lives at data/superclass_mapping_known.csv, that's what the makefile targets use by default:

Class,Superclass,Class ID,Superclass ID
tts_models/en/ljspeech/fast_pitch,FeedForward,0,0
tts_models/en/ljspeech/speedy-speech,FeedForward,1,0
tts_models/en/ljspeech/vits--neon,VITS,2,1

Superclass IDs of -1 are auto-reassigned to unique buckets by the loader for "unique architecture" entries that don't share a group.

H-Shared — Shared Hierarchical Classifier

uv run scripts/training/train_hier.py \
  --path_to_dataset "data/prepared_ds_seg_enc" \
  --pre_encoded True --pre_augmented True --is_segmented True \
  --out_folder "exp/trained_models/hier_mt_arc_m03" \
  --use_arc_margin --easy_margin True --arc_m 0.3 \
  --weighted_sampling True \
  --hierarchy_type "H-Shared" \
  --superclass_lut "data/superclass_mapping_known.csv"

Or:

make train_hier_shared

H-Arch — Architecture-Specific Hierarchical Classifier

uv run scripts/training/train_hier.py \
  --path_to_dataset "data/prepared_ds_seg_enc" \
  --pre_encoded True --pre_augmented True --is_segmented True \
  --out_folder "exp/trained_models/hier_cas_arc_m03" \
  --use_arc_margin --easy_margin True --arc_m 0.3 \
  --weighted_sampling True \
  --hierarchy_type "H-Arch" \
  --superclass_lut "data/superclass_mapping_known.csv"

Or:

make train_hier_arch

Key arguments (both scripts)

Argument Default Notes
--num_classes 24 Set to your ID-class count
--num_epochs 30
--batch_size 256 Drop if VRAM-bound
--feat_dim 768 Matches wav2vec2-base hidden size
--lr 1e-3
--interval 10 LR decay interval (epochs)
--seed 688
--weighted_sampling False Recommended ON for the imbalanced MLAAD training split
--use_arc_margin flag Mutually informative with --use_sub-center-arc_margin
--arc_m 0.5 Margin
--arc_s 30 Scale (or "auto")
--easy_margin False
--resume-checkpoint None Path to anti-spoofing_feat_model.pth to resume from
--resume-epoch None Pair with --resume-checkpoint

Flat-only:

  • --use_center_loss, --center_loss_weight (joint center loss)

Hierarchical-only:

  • --hierarchy_type ∈ {H-Shared, H-Arch} (required)
  • --superclass_lut <path> (accepts a repo-relative path like data/superclass_mapping_known.csv or a name relative to --path_to_dataset)

Outputs

exp/trained_models/<run_name>/
├── anti-spoofing_feat_model.pth   # final weights (loaded by eval scripts)
├── args.json                       # full training arg snapshot
├── training_stats.json             # per-epoch loss/accuracy
├── optimizer.pth
├── checkpoint/                     # per-epoch snapshots (if save_interval hits)
└── logs/                           # TensorBoard event files

args.json is what get_classification_metrics.py and ood_detector.py read to know what model to instantiate (flat vs H-Shared vs H-Arch, ArcFace vs sub-center, etc.).

Inspect TensorBoard logs:

uv run tensorboard --logdir exp/trained_models/<run_name>/logs

Multi-seed runs

Driver script: scripts/run_multiseed_ci.py. It runs train_hier.py for a list of seeds and aggregates final metrics.

uv run scripts/run_multiseed_ci.py \
  --workspace . \
  --train_dataset data/prepared_ds_seg_enc \
  --eval_dataset data/prepared_ds_seg_enc \
  --seeds 688,699,710 \
  --base_out exp/multiseed \
  --num_epochs 30 --batch_size 256

Use --dry_run to print the commands without launching them.

Backbone ablation

Driver script: scripts/run_backbone_ablation.py — re-encodes the dataset with alternative SSL backbones (HuBERT, etc.) and trains the same hierarchical head for each.

uv run scripts/run_backbone_ablation.py \
  --workspace . \
  --raw_dataset data/MLAADv5_for_sourcetracing \
  --train_dataset data/prepared_ds_seg_enc \
  --eval_dataset data/prepared_ds_seg_enc \
  --out_root exp/backbone_ablation \
  --backbones w2v2,hubert \
  --run

Drop --run to print-only.